#!/usr/bin/env python3
"""
从 ai-planning 组装 SFT 训练/验证集，并启动 zk_trainer SFT 训练。

用法（先 cd "$AUTORESEARCH_CHAT_ROOT"，所有产物落 chat 工作区）:
  # 仅组装数据
  python scripts/prepare_and_train_sft.py prepare --output_dir "$AUTORESEARCH_CHAT_ROOT/sft_data"

  # 启动训练（需先 prepare）
  python scripts/prepare_and_train_sft.py train \
    --data_dir "$AUTORESEARCH_CHAT_ROOT/sft_data" \
    --model_path /mnt/wangsenhao/verl_zk/Qwen3-4B-Instruct-2507 \
    --model_type qwen3 \
    --train_output "$AUTORESEARCH_CHAT_ROOT/sft_output"
"""

import argparse
import json
import logging
import os
import shutil
import subprocess
import sys
from pathlib import Path

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
log = logging.getLogger(__name__)

CHAT_ROOT = Path(__file__).resolve().parent.parent
AI_PLANNING = CHAT_ROOT / "ai-planning"
TRAIN_SET_DIR = AI_PLANNING / "data" / "train_set"
DATA_TRAIN_DIR = AI_PLANNING / "data_train"
GENERATE_TRAIN_SCRIPT = DATA_TRAIN_DIR / "generate_train.py"


# ──────────────────────────────────────────────
# 1. 数据组装
# ──────────────────────────────────────────────

def regen_jsonl_from_csv():
    """调用 data_train/generate_train.py 从 CSV 重新生成 all_train.jsonl。"""
    if not GENERATE_TRAIN_SCRIPT.exists():
        log.error(f"generate_train.py 不存在: {GENERATE_TRAIN_SCRIPT}")
        return False
    log.info("从 CSV 重新生成 JSONL ...")
    result = subprocess.run(
        [sys.executable, str(GENERATE_TRAIN_SCRIPT)],
        cwd=str(DATA_TRAIN_DIR),
    )
    if result.returncode != 0:
        log.error("generate_train.py 执行失败")
        return False
    log.info("CSV → JSONL 完成")
    return True


def scan_train_set():
    """扫描 train_set/ 下所有子目录，返回 (train_files, valid_files)。"""
    train_files, valid_files = [], []
    if not TRAIN_SET_DIR.exists():
        log.error(f"train_set 目录不存在: {TRAIN_SET_DIR}")
        return train_files, valid_files

    for subdir in sorted(TRAIN_SET_DIR.iterdir()):
        if not subdir.is_dir() or subdir.name.startswith("."):
            continue
        for f in sorted(subdir.iterdir()):
            if not f.is_file() or f.suffix != ".jsonl":
                continue
            if "_valid" in f.name:
                valid_files.append(f)
            else:
                train_files.append(f)
    return train_files, valid_files


def count_lines(path):
    with open(path, "r", encoding="utf-8") as f:
        return sum(1 for _ in f)


def prepare_data(output_dir: str, regen: bool = False, dataset_name: str = "zk_sft_new_structure"):
    """
    组装训练集和验证集到 output_dir，生成 zk_trainer 所需的目录结构:
      output_dir/
        train/<dataset_name>/merged_train.jsonl
        validation/<dataset_name>/part-0.jsonl
    """
    output = Path(output_dir).resolve()
    train_dir = output / "train" / dataset_name
    valid_dir = output / "validation" / dataset_name

    if regen:
        regen_jsonl_from_csv()

    train_files, valid_files = scan_train_set()
    if not train_files:
        log.error("未找到训练文件")
        return None

    log.info(f"训练文件 ({len(train_files)}):")
    for f in train_files:
        log.info(f"  {f.relative_to(AI_PLANNING)}  ({count_lines(f)} 条)")
    log.info(f"验证文件 ({len(valid_files)}):")
    for f in valid_files:
        log.info(f"  {f.relative_to(AI_PLANNING)}  ({count_lines(f)} 条)")

    # 清理旧数据
    for d in [train_dir, valid_dir]:
        if d.exists():
            shutil.rmtree(d)
        d.mkdir(parents=True)

    # 合并训练文件
    merged_path = train_dir / "merged_train.jsonl"
    total_train = 0
    with open(merged_path, "w", encoding="utf-8") as out:
        for fp in train_files:
            with open(fp, "r", encoding="utf-8") as inp:
                for line in inp:
                    out.write(line)
                    total_train += 1
    log.info(f"合并训练集: {total_train} 条 → {merged_path}")

    # 合并验证文件
    merged_valid_path = valid_dir / "part-0.jsonl"
    total_valid = 0
    with open(merged_valid_path, "w", encoding="utf-8") as out:
        for fp in valid_files:
            with open(fp, "r", encoding="utf-8") as inp:
                for line in inp:
                    out.write(line)
                    total_valid += 1
    log.info(f"验证集: {total_valid} 条 → {merged_valid_path}")

    log.info("=" * 50)
    log.info(f"数据目录: {output}")
    log.info(f"  train:      {total_train} 条")
    log.info(f"  validation: {total_valid} 条")

    return str(output)


# ──────────────────────────────────────────────
# 2. SFT 训练
# ──────────────────────────────────────────────

def run_sft_training(data_dir: str, model_path: str, model_type: str, train_output: str,
                     epochs: int = 3, lr: float = 1e-5, max_seq_length: int = 1024,
                     per_device_batch: int = 1, grad_accum: int = 1):
    """基于 zk_trainer 的 accelerate launch 方式启动 SFT 训练。"""
    zk_trainer_dir = str(CHAT_ROOT / "zk_trainer")
    if not os.path.isdir(zk_trainer_dir):
        log.error(f"zk_trainer 目录不存在: {zk_trainer_dir}")
        return None

    data_dir = str(Path(data_dir).resolve())
    train_output = str(Path(train_output).resolve())
    os.makedirs(train_output, exist_ok=True)

    run_name = f"{model_type}-zk"

    env = os.environ.copy()
    env["PYTHONPATH"] = f"{zk_trainer_dir}:{env.get('PYTHONPATH', '')}"
    env["WANDB_DISABLED"] = "true"
    env["WANDB_MODE"] = "offline"
    env.setdefault("VOLUME_PREFIX", "/mnt/xiaoai-zk-model-train-tj5")

    # Step 1: 用 get_accelerate_config.py 动态生成 accelerate config
    num_gpu_result = subprocess.run(
        ["nvidia-smi", "--query-gpu=gpu_name", "--format=csv,noheader"],
        capture_output=True, text=True
    )
    num_processes = len(num_gpu_result.stdout.strip().split("\n")) if num_gpu_result.returncode == 0 else 1
    log.info(f"检测到 {num_processes} 张 GPU")

    acc_config_dir = zk_trainer_dir
    gen_config_cmd = [
        sys.executable, f"{zk_trainer_dir}/llm/train/get_accelerate_config.py",
        "--dt_framework=fsdp",
        "--fsdp_sharding_strategy=1",
        "--machine_rank=0",
        "--num_machines=1",
        f"--num_processes={num_processes}",
        f"--output_dir={acc_config_dir}",
        "--main_process_ip=127.0.0.1",
        "--main_process_port=56390",
        "--mixed_precision=bf16",
        f"--model_type={model_type}",
        "--ds_stage=3",
        "--fsdp_version=2",
        "--fsdp_wrap_cls=",
    ]
    log.info("生成 accelerate config ...")
    result = subprocess.run(gen_config_cmd, env=env, cwd=zk_trainer_dir)
    if result.returncode != 0:
        log.error("生成 accelerate config 失败")
        return None

    acc_config = os.path.join(acc_config_dir, "accelerate_config_0.yaml")
    if not os.path.exists(acc_config):
        log.error(f"accelerate config 未生成: {acc_config}")
        return None
    log.info(f"accelerate config: {acc_config}")

    # Step 2: 启动训练
    train_cmd = [
        "accelerate", "launch",
        "--main_process_port", "12345",
        "--config_file", acc_config,
        f"{zk_trainer_dir}/llm/train/trainer_general.py",
        "--model_name_or_path", model_path,
        "--model_type", model_type,
        "--use_auto_class", "true",
        "--output_dir", train_output,
        "--dataset_type", "zk_sft",
        "--train_data_dir", os.path.join(data_dir, "train"),
        "--valid_data_dir", os.path.join(data_dir, "validation"),
        "--do_eval",
        "--fp16_full_eval",
        "--per_device_eval_batch_size", "1",
        "--include_for_metrics", "inputs",
        "--batch_eval_metrics",
        "--concat_samples", "False",
        "--mix_at_eval", "False",
        "--num_train_epochs", str(epochs),
        "--learning_rate", str(lr),
        "--gradient_accumulation_steps", str(grad_accum),
        "--ignore_data_skip", "False",
        "--resume_train_if_ckpt_exists", "false",
        "--save_strategy", "no",
        "--logging_steps", "5",
        "--save_total_limit", "0",
        "--max_seq_length", str(max_seq_length),
        "--adam_beta1", "0.9",
        "--adam_beta2", "0.95",
        "--adam_epsilon", "1e-9",
        "--warmup_ratio", "0.1",
        "--lr_scheduler_type", "cosine",
        "--use_shuffle", "true",
        "--batch_eval_metrics",
        "--run_name", run_name,
        "--report_to", "none",
        "--bf16",
    ]

    log.info("=" * 50)
    log.info("启动 SFT 训练")
    log.info(f"  模型:   {model_path}")
    log.info(f"  数据:   {data_dir}")
    log.info(f"  输出:   {train_output}")
    log.info(f"  GPU:    {num_processes}")
    log.info(f"  epochs: {epochs}, lr: {lr}, max_seq: {max_seq_length}")

    result = subprocess.run(train_cmd, env=env, cwd=zk_trainer_dir)

    if result.returncode != 0:
        log.error(f"训练失败，返回码: {result.returncode}")
        return None

    log.info(f"训练完成，输出: {train_output}")
    return train_output


# ──────────────────────────────────────────────
# CLI
# ──────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(description="ai-planning SFT 数据组装 & 训练")
    subparsers = parser.add_subparsers(dest="command")

    # prepare
    p_prepare = subparsers.add_parser("prepare", help="仅组装数据")
    p_prepare.add_argument("--output_dir", required=True, help="数据输出目录")
    p_prepare.add_argument("--regen_from_csv", action="store_true", help="从 CSV 重新生成 JSONL")
    p_prepare.add_argument("--dataset_name", default="zk_sft_new_structure")

    # train
    p_train = subparsers.add_parser("train", help="启动训练（需先 prepare）")
    p_train.add_argument("--data_dir", required=True, help="已组装的数据目录（prepare 的 output_dir）")
    p_train.add_argument("--model_path", required=True, help="基模路径")
    p_train.add_argument("--model_type", default="qwen3", help="模型类型")
    p_train.add_argument("--train_output", required=True, help="训练输出目录")
    p_train.add_argument("--epochs", type=int, default=3)
    p_train.add_argument("--lr", type=float, default=1e-5)
    p_train.add_argument("--max_seq_length", type=int, default=1024)

    args = parser.parse_args()

    if args.command == "prepare":
        prepare_data(args.output_dir, regen=args.regen_from_csv, dataset_name=args.dataset_name)

    elif args.command == "train":
        run_sft_training(
            data_dir=args.data_dir,
            model_path=args.model_path,
            model_type=args.model_type,
            train_output=args.train_output,
            epochs=args.epochs,
            lr=args.lr,
            max_seq_length=args.max_seq_length,
        )
    else:
        parser.print_help()


if __name__ == "__main__":
    main()