"""
真实训练执行器（基于项目内 train.py 调用 LLaMA-Factory 库）

- 根据任务配置构建 `python train.py` 命令（非 llamafactory-cli 子命令）
- 实时捕获训练日志与 trainer_log.jsonl 的 loss，回写到 PlatformStore
- 支持 SIGSTOP / SIGCONT / SIGKILL 实现暂停 / 恢复 / 取消
- 训练完成后自动绘制 loss 曲线（matplotlib 可选）
"""
from __future__ import annotations

import json
import os
import queue
import signal
import subprocess
import sys
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Any

from app.db.platform_store import get_platform_store

BACKEND_ROOT = Path(__file__).resolve().parent.parent.parent.parent
DATA_DIR = BACKEND_ROOT / "data"
DATASET_INFO_PATH = DATA_DIR / "dataset_info.json"
DATASET_STORE_DIR = DATA_DIR / "fine_tune_datasets"
OUTPUT_ROOT = DATA_DIR / "fine_tune_outputs"
TRAIN_SCRIPT = BACKEND_ROOT / "train.py"

_running_processes: dict[str, "subprocess.Popen[Any]"] = {}

_stage_map = {"sft": "sft", "dpo": "dpo", "cpt": "pt", "cot": "sft"}


# ──────────────────────────────────────────────
# 工具函数
# ──────────────────────────────────────────────
def _log(store: Any, task_id: str, msg: str, log_type: str = "info") -> None:
    try:
        task = store.task(task_id)
        logs = list(task.get("logs") or [])
        logs.append({"time": datetime.now().strftime("%H:%M:%S"), "msg": msg, "type": log_type})
        store.update_task_runtime(task_id, extra={"logs": logs})
    except Exception:
        pass


def _resolve_model_path(store: Any, name_or_path: str) -> str:
    """把模型短名解析为本地绝对路径；已是合法路径则原样返回。"""
    if not name_or_path:
        return name_or_path
    p = Path(name_or_path)
    if p.exists() and (p / "config.json").exists():
        return str(p.resolve())
    for m in store.models() or []:
        if m.get("name") == name_or_path or m.get("path") == name_or_path:
            resolved = Path(m.get("path", ""))
            if resolved.exists():
                return str(resolved.resolve())
    return name_or_path


def _materialize_dataset(store: Any, dataset_id: str) -> str:
    """把数据集管理系统的 UUID 数据集落盘并注册进 dataset_info.json，返回 --dataset key。"""
    if not dataset_id or dataset_id == "identity":
        return dataset_id
    try:
        existing = json.loads(DATASET_INFO_PATH.read_text(encoding="utf-8")) if DATASET_INFO_PATH.exists() else {}
    except Exception:
        existing = {}
    if dataset_id in existing:
        return dataset_id
    try:
        ds = store.dataset(dataset_id)
    except Exception:
        return dataset_id
    files = ds.get("files") or []
    if not files:
        return dataset_id
    file_id = files[0].get("id")
    ext = files[0].get("ext", ".json")
    if ext not in (".json", ".jsonl"):
        ext = ".json"
    if not file_id:
        return dataset_id
    try:
        row = store.dataset_file(file_id)
    except Exception:
        return dataset_id
    content = row.get("content", "")
    DATASET_STORE_DIR.mkdir(parents=True, exist_ok=True)
    actual = DATASET_STORE_DIR / f"{dataset_id}{ext}"
    actual.write_text(content, encoding="utf-8")
    rel = actual.relative_to(DATA_DIR)
    entry: dict[str, Any] = {"file_name": str(rel)}
    try:
        text = content.strip()
        if text.startswith("["):
            text = text[text.find("{") : text.find("}") + 1]
        sample = json.loads(text)
        if isinstance(sample, dict) and ("messages" in sample or "conversations" in sample):
            key = "messages" if "messages" in sample else "conversations"
            entry["formatting"] = "sharegpt"
            entry["columns"] = {"messages": key}
    except Exception:
        pass
    existing[dataset_id] = entry
    DATASET_INFO_PATH.parent.mkdir(parents=True, exist_ok=True)
    DATASET_INFO_PATH.write_text(json.dumps(existing, ensure_ascii=False, indent=2), encoding="utf-8")
    return dataset_id


# ──────────────────────────────────────────────
# 训练主流程
# ──────────────────────────────────────────────
def run_training(task_id: str) -> None:
    store = get_platform_store()
    try:
        task = store.task(task_id)
    except Exception:
        return

    cfg = dict(task)
    mode = (cfg.get("train_type") or "sft").lower()
    stage = _stage_map.get(mode, "sft")

    base_model = _resolve_model_path(store, cfg.get("base_model", ""))
    train_dataset = _materialize_dataset(store, cfg.get("train_dataset_id", ""))
    eval_dataset = _materialize_dataset(store, cfg.get("eval_dataset_id") or cfg.get("eval_dataset", ""))

    finetuning_type = (cfg.get("train_method") or "lora").lower()
    OUTPUT_ROOT.mkdir(parents=True, exist_ok=True)
    output_dir = str(OUTPUT_ROOT / task["name"])

    env = os.environ.copy()
    env.setdefault("HF_HUB_OFFLINE", "1")
    env.setdefault("TRANSFORMERS_OFFLINE", "1")
    env.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

    gpus = cfg.get("gpus") or [0]
    num_gpus = int(cfg.get("num_gpus", len(gpus)) or 1) or 1
    # 仅当显式指定非默认 GPU 时限制可见设备；多卡统一走 torchrun
    if gpus and str(gpus[0]) not in ("0", "gpu-0"):
        env["CUDA_VISIBLE_DEVICES"] = ",".join(str(g).replace("gpu-", "") for g in gpus)
        num_gpus = len(gpus)

    if num_gpus > 1:
        _log(store, task_id, f"[INFO] 启用分布式训练: {num_gpus} 个 GPU", "info")
        cmd = [sys.executable, "-m", "torch.distributed.run", "--nproc_per_node", str(num_gpus), str(TRAIN_SCRIPT)]
    else:
        cmd = [sys.executable, str(TRAIN_SCRIPT)]
    cmd += [
        "--stage", stage,
        "--do_train", "True",
        "--model_name_or_path", base_model,
        "--dataset", train_dataset,
        "--dataset_dir", str(DATA_DIR),
        "--template", cfg.get("template", "default"),
        "--finetuning_type", finetuning_type,
        "--output_dir", output_dir,
        "--trust_remote_code", "True",
        "--overwrite_output_dir", "True",
        "--report_to", "none",
        "--learning_rate", str(cfg.get("learning_rate", "1e-5")),
        "--num_train_epochs", str(cfg.get("n_epochs", 3)),
        "--per_device_train_batch_size", str(cfg.get("batch_size", 4)),
        "--cutoff_len", str(cfg.get("max_length", 1024)),
        "--gradient_accumulation_steps", str(cfg.get("gradient_accumulation_steps", 4)),
        "--max_samples", str(cfg.get("max_samples", 100000)),
        "--lr_scheduler_type", cfg.get("lr_scheduler_type", "cosine"),
        "--warmup_ratio", str(cfg.get("warmup_ratio", 0.03)),
        "--max_grad_norm", str(cfg.get("max_grad_norm", "1.0")),
        "--optim", cfg.get("optim", "adamw_torch"),
        "--logging_steps", str(cfg.get("logging_steps", 10)),
        "--save_steps", str(cfg.get("save_steps", 100)),
        "--save_total_limit", str(cfg.get("save_total_limit", 5)),
        "--flash_attn", cfg.get("flash_attn", "auto"),
    ]
    dtype = cfg.get("dtype", "bf16")
    if dtype == "bf16":
        cmd += ["--bf16", "True"]
    elif dtype == "fp16":
        cmd += ["--fp16", "True"]
    if finetuning_type == "lora":
        cmd += [
            "--lora_rank", str(cfg.get("lora_rank", 8)),
            "--lora_alpha", str(cfg.get("lora_alpha", 16)),
            "--lora_dropout", str(cfg.get("lora_dropout", "0.05")),
            "--lora_target", cfg.get("lora_target", "all"),
        ]
    if cfg.get("do_eval", False):
        cmd += ["--do_eval", "True", "--eval_strategy", "steps",
                "--eval_steps", str(cfg.get("eval_steps", 50)),
                "--per_device_eval_batch_size", "4"]
        if eval_dataset:
            cmd += ["--eval_dataset", eval_dataset]
        else:
            cmd += ["--val_size", str(cfg.get("val_size", "0.1"))]
    if cfg.get("resume_from"):
        cmd += ["--resume_from_checkpoint", str(cfg["resume_from"])]

    _log(store, task_id, f"[INFO] 训练任务启动: {mode.upper()} 微调")
    _log(store, task_id, f"[INFO] 基座模型: {base_model}")
    _log(store, task_id, f"[INFO] 数据集: {train_dataset}")
    _log(store, task_id, f"[INFO] 输出目录: {output_dir}")
    store.update_task_runtime(task_id, status="running", progress=10, process_id=None)

    try:
        process = subprocess.Popen(
            cmd,
            cwd=str(BACKEND_ROOT),
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT,
            text=True,
            env=env,
            start_new_session=True,
        )
        _running_processes[task_id] = process
        store.update_task_runtime(task_id, process_id=process.pid)
        _log(store, task_id, f"[INFO] 训练进程已启动 (PID: {process.pid})", "info")

        log_file = Path(output_dir) / "trainer_log.jsonl"
        last_pos = 0
        stop_flag = threading.Event()
        finished = threading.Event()

        def watch_logs() -> None:
            nonlocal last_pos
            while not stop_flag.is_set():
                time.sleep(1)
                if not log_file.exists():
                    continue
                try:
                    with open(log_file, "r", encoding="utf-8") as lf:
                        lf.seek(last_pos)
                        for line in lf:
                            try:
                                data = json.loads(line)
                                if "loss" not in data:
                                    continue
                                loss = float(data["loss"])
                                history = list(store.task(task_id).get("loss_history") or [])
                                history.append(loss)
                                extra: dict[str, Any] = {"current_loss": loss, "loss_history": history}
                                if data.get("lr") is not None:
                                    extra["learning_rate"] = float(str(data["lr"]).replace("'", ""))
                                if data.get("epoch") is not None:
                                    extra["current_epoch"] = float(data["epoch"])
                                if data.get("percentage") is not None:
                                    extra["progress"] = float(data["percentage"])
                                if data.get("remaining_time"):
                                    extra["eta"] = str(data["remaining_time"])
                                store.update_task_runtime(task_id, extra=extra)
                                step = data.get("current_steps")
                                total = data.get("total_steps")
                                _log(store, task_id, f"Step {step}/{total} loss={loss:.4f}")
                                try:
                                    if int(step) >= int(total):
                                        finished.set()
                                except Exception:
                                    pass
                            except Exception:
                                pass
                        last_pos = lf.tell()
                except Exception:
                    pass

        watcher = threading.Thread(target=watch_logs, daemon=True)
        watcher.start()

        q: "queue.Queue[str]" = queue.Queue()

        def reader() -> None:
            try:
                for line in process.stdout:
                    q.put(line.rstrip("\r\n"))
            except Exception:
                pass
            finally:
                q.put("")

        r = threading.Thread(target=reader, daemon=True)
        r.start()

        while True:
            try:
                raw = q.get(timeout=1)
            except queue.Empty:
                if finished.is_set() or (process.poll() is not None and q.empty()):
                    try:
                        raw = q.get(timeout=1)
                    except queue.Empty:
                        break
                else:
                    continue
            if not raw:
                break
            low = raw.lower()
            log_type = "error" if "error" in low else ("warn" if "warn" in low else "info")
            _log(store, task_id, raw, log_type)

        stop_flag.set()
        try:
            process.stdout.close()
        except Exception:
            pass
        watcher.join(timeout=3)
        try:
            process.wait(timeout=60)
        except subprocess.TimeoutExpired:
            process.kill()
            process.wait()

        ret = process.returncode
        if ret == 0:
            store.update_task_runtime(task_id, status="completed", progress=100)
            _log(store, task_id, "[INFO] 训练完成！", "info")
            try:
                store.ensure_trained_model_for_task(task_id, cfg, output_dir)
            except Exception as exc:
                _log(store, task_id, f"[WARN] 登记训练产物失败: {exc}", "warn")
            try:
                _plot_loss_curve(task_id, output_dir, bool(cfg.get("do_eval", False)))
            except Exception as exc:
                _log(store, task_id, f"[WARN] Loss 曲线异常: {exc}", "warn")
        else:
            store.update_task_runtime(task_id, status="failed", extra={"error_message": f"训练退出码: {ret}"})
            _log(store, task_id, f"[ERROR] 训练退出码: {ret}", "error")
    except FileNotFoundError:
        store.update_task_runtime(task_id, status="failed", extra={"error_message": "未找到 train.py 或 Python，请确认后端根目录存在 train.py 且 LLaMA-Factory 已安装"})
        _log(store, task_id, "[ERROR] 未找到 train.py / Python，无法启动真实训练", "error")
    except Exception as exc:  # noqa: BLE001
        store.update_task_runtime(task_id, status="failed", extra={"error_message": str(exc)})
        _log(store, task_id, f"[ERROR] 训练异常: {exc}", "error")
    finally:
        _running_processes.pop(task_id, None)


# ──────────────────────────────────────────────
# 进程信号控制
# ──────────────────────────────────────────────
def _signal(task_id: str, sig: int) -> bool:
    process = _running_processes.get(task_id)
    if not process:
        return False
    try:
        try:
            pgid = os.getpgid(process.pid)
            os.killpg(pgid, sig)
        except (ProcessLookupError, PermissionError):
            process.send_signal(sig)
        return True
    except Exception:
        return False


def pause(task_id: str) -> bool:
    ok = _signal(task_id, signal.SIGSTOP)
    if ok:
        get_platform_store().update_task_runtime(task_id, status="paused")
        _log(get_platform_store(), task_id, "[INFO] 训练已暂停", "info")
    return ok


def resume(task_id: str) -> bool:
    ok = _signal(task_id, signal.SIGCONT)
    if ok:
        get_platform_store().update_task_runtime(task_id, status="running")
        _log(get_platform_store(), task_id, "[INFO] 训练已继续", "info")
    return ok


def cancel(task_id: str) -> bool:
    ok = _signal(task_id, signal.SIGKILL)
    if ok:
        get_platform_store().update_task_runtime(task_id, status="failed", extra={"error_message": "用户中断训练"})
        _log(get_platform_store(), task_id, "[WARN] 用户中断了训练", "warn")
    return ok


def _plot_loss_curve(task_id: str, output_dir: str, has_eval: bool = False) -> None:
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        state_file = Path(output_dir) / "trainer_state.json"
        if not state_file.exists():
            return
        state = json.loads(state_file.read_text(encoding="utf-8"))
        log_history = state.get("log_history", [])
        if not log_history:
            return
        train_steps, train_losses, eval_steps, eval_losses = [], [], [], []
        for entry in log_history:
            if "loss" in entry and "step" in entry:
                train_steps.append(entry["step"])
                train_losses.append(entry["loss"])
            if has_eval and "eval_loss" in entry and "step" in entry:
                eval_steps.append(entry["step"])
                eval_losses.append(entry["eval_loss"])
        if not train_losses:
            return
        fig, ax = plt.subplots(figsize=(10, 5))
        ax.plot(train_steps, train_losses, label="Training Loss", color="#409eff", linewidth=1.5)
        ax.axhline(y=min(train_losses), color="#67c23a", linestyle="--", alpha=0.5,
                   label=f"Min: {min(train_losses):.4f}")
        if eval_losses:
            ax.plot(eval_steps, eval_losses, label="Validation Loss", color="#f56c6c",
                    linewidth=1.5, marker="o", markersize=3)
        ax.set_xlabel("Step")
        ax.set_ylabel("Loss")
        ax.set_title("Training Loss Curve")
        ax.legend()
        ax.grid(True, alpha=0.3)
        save_path = Path(output_dir) / "loss_curve.png"
        fig.savefig(str(save_path), dpi=150, bbox_inches="tight")
        plt.close(fig)
        _log(get_platform_store(), task_id, f"Loss 曲线已保存: {save_path}")
    except Exception:
        pass