#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""训练任务脚本（Go 后端 --task-json 驱动，产物契约见 server/common/training_runner.go）。

用法：
    python train_server.py --task-json tasks/<taskId>.json

任务参数（Go 侧写入，字段相对训练机 workdir）：
    workdir       训练机工作目录（脚本 / yolov8n.pt / venv 所在），启动即 chdir
    yolo          训练集目录（含 dataset.yaml），相对 workdir
    imgsz         训练/导出分辨率（默认 704，与端侧推理对齐）
    epochs        训练轮数
    batch         批大小
    device        GPU 编号或 cpu
    project       训练输出目录（相对 workdir，末尾自动拼 name）
    log_file      每 epoch 一行 JSON 的进度日志（相对 workdir）
    result_file   结束结果 JSON（相对 workdir）

产物契约：
    log_file      {"epoch":1,"total":150,"metrics":{"metrics/mAP50(B)":0.87,...}}
    result_file   {"metrics":{...},"names":["pheasant","suspect"],"best_tflite":"runs/tasks/<id>/weights/best.tflite",
                   "tflite_check":{"ok":true,"reason":"","inputs":[...],"outputs":[...]}}
                  result_file 存在 = 训练完成；异常时写 {"error":"..."}，Go 侧据以置失败并展示原因。
                  tflite_check.ok=false（产物 shape 异常）时 Go 侧置训练失败并带出 reason。
"""

import argparse
import json
import math
import os
import struct
import sys
import traceback
from pathlib import Path

# 抗尺度漂移（与人工训练基线一致）：随机缩放输入 0.5~1.5x
MULTI_SCALE = 0.5
# 早停耐心（连续 N 轮无提升即停）
PATIENCE = 30


def on_fit_epoch_end(trainer):
    """ultralytics 回调：每 epoch 结束写一行进度 JSON（行缓冲，进程被杀不丢行）。"""
    metrics = {}
    for k, v in (trainer.metrics or {}).items():
        if isinstance(v, (int, float)) and math.isfinite(v):
            metrics[k] = round(float(v), 5)
    if _LOG is None:
        return
    total = getattr(trainer, "epochs", None) or getattr(trainer, "total_epochs", 0) or 0
    _LOG.write(json.dumps({
        "epoch": trainer.epoch + 1,
        "total": total,
        "metrics": metrics,
    }, ensure_ascii=False) + "\n")


def register_callback(model):
    # ultralytics ≥8.4 中 trainer 使用 model 实例的 callbacks 字典（model.train 前以
    # _callbacks=self.callbacks 传入），模块级 get_default_callbacks() 的 append 不生效；
    # add_callback 追加到 model 字典即随训练触发，进度 jsonl 才有输出
    model.add_callback("on_fit_epoch_end", on_fit_epoch_end)


def write_result(result_file, payload):
    Path(result_file).parent.mkdir(parents=True, exist_ok=True)
    # tmp + rename 原子写，避免 Go 侧读到半截文件
    tmp = result_file + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(payload, f, ensure_ascii=False, indent=2)
    os.replace(tmp, result_file)


def read_names(yolo_dir):
    """从 data.yaml 读类别名（有序列表，ultralytics 依赖 pyyaml）"""
    try:
        import yaml
        with open(Path(yolo_dir) / "dataset.yaml", encoding="utf-8") as f:
            cfg = yaml.safe_load(f)
        names = cfg.get("names") or {}
        if isinstance(names, dict):
            return [names[k] for k in sorted(names, key=lambda k: int(k))]
        return list(names)
    except Exception:
        return []


def find_best_tflite(save_dir):
    """定位导出的 tflite（ultralytics 各版本产物位置不一）：
    优先 fp32 best.tflite（与人工基线一致），再 float32/int8，最后兜底全局搜。"""
    weights = save_dir / "weights"
    for name in ("best.tflite", "best_float32.tflite", "best_int8.tflite"):
        p = weights / name
        if p.exists():
            return p
    candidates = sorted(weights.rglob("*.tflite"))
    if candidates:
        return candidates[0]
    return None


# TFLite flatbuffer 张量类型（tensorflow/lite/schema/schema.fbs，字段顺序敏感）
TENSOR_TYPE = {0: "FLOAT32", 1: "FLOAT16", 2: "INT32", 3: "UINT8", 4: "INT64",
               5: "STRING", 6: "BOOL", 7: "INT16", 8: "COMPLEX64", 9: "INT8",
               10: "FLOAT64", 11: "COMPLEX128", 12: "UINT64", 13: "RESOURCE",
               14: "VARIANT", 15: "UINT32", 16: "UINT16", 17: "INT4",
               18: "BFLOAT16", 19: "FLOAT8_E4M3FN", 20: "FLOAT8_E4M3FNUZ",
               21: "FLOAT8_E5M2", 22: "FLOAT8_E5M2FNUZ"}


class _FB:
    """极简 flatbuffer 读取器（仅 Model/SubGraph/Tensor 表所需字段）"""

    def __init__(self, data):
        self.d = data

    def u32(self, pos): return struct.unpack_from("<I", self.d, pos)[0]
    def i32(self, pos): return struct.unpack_from("<i", self.d, pos)[0]
    def u16(self, pos): return struct.unpack_from("<H", self.d, pos)[0]

    def field(self, t, i):
        vto = self.i32(t)
        vt = t - vto
        vs = self.u16(vt)
        off = vt + 4 + 2 * i
        if off + 2 > vt + vs:
            return None
        f = self.u16(off)
        return None if f == 0 else t + f

    def deref(self, pos):
        return pos + self.u32(pos) if pos is not None else None

    def vec(self, pos):
        pos = self.deref(pos)
        if pos is None:
            return None
        return pos + 4, self.u32(pos)

    def vec_table(self, pos):
        start, n = self.vec(pos)
        return [start + i * 4 + self.u32(start + i * 4) for i in range(n)]

    def string(self, pos):
        pos = self.deref(pos)
        if pos is None:
            return None
        n = self.u32(pos)
        return self.d[pos + 4:pos + 4 + n].decode("utf-8", "replace")

    def int_vec(self, pos):
        start, n = self.vec(pos)
        return [self.i32(start + 4 * i) for i in range(n)]


def check_tflite(path, imgsz):
    """TFLite 产物自检（原 inspect_tflite.py，2026-08-26 并入）：
    输入恰 1 张且 4 维、元素总数 == imgsz²×3（NCHW/NHWC 皆可）、输出 ≥ 1 张且 batch 维 = 1。
    返回 {"ok","reason","inputs":[{"name","shape","type"}],"outputs":[...]}。"""
    try:
        with open(path, "rb") as f:
            fb = _FB(f.read())
        root = fb.u32(0)                          # Model 表
        subgraphs = fb.vec_table(fb.field(root, 2))  # Model.subgraphs [2]
        sg = subgraphs[0]
        tensors = fb.vec_table(fb.field(sg, 0))   # SubGraph.tensors [0]
        inputs = fb.int_vec(fb.field(sg, 1))      # SubGraph.inputs [1]
        outputs = fb.int_vec(fb.field(sg, 2))     # SubGraph.outputs [2]

        def describe(i):
            t = tensors[i]
            shape = fb.int_vec(fb.field(t, 0))
            # type 有默认值 FLOAT32(=0)：字段等于默认值时 flatbuffer 省略（vtable 偏移 0）
            ttype_f = fb.field(t, 1)
            ttype = fb.d[ttype_f] if ttype_f else 0
            return {"name": fb.string(fb.field(t, 3)) or "?",
                    "shape": shape,
                    "type": TENSOR_TYPE.get(ttype, ttype)}

        in_desc = [describe(i) for i in inputs]
        out_desc = [describe(i) for i in outputs]

        problems = []
        if len(in_desc) != 1:
            problems.append(f"输入张量 {len(in_desc)} 个（期望 1）")
        else:
            shape = in_desc[0]["shape"]
            n = 1
            for d in shape:
                n *= d
            if len(shape) != 4:
                problems.append(f"输入 shape {shape} 非 4 维")
            elif n != imgsz * imgsz * 3:
                problems.append(f"输入元素数 {n} ≠ imgsz²×3={imgsz * imgsz * 3}（shape {shape}，imgsz={imgsz}）")
        if not out_desc:
            problems.append("无输出张量")
        else:
            for o in out_desc:
                if o["shape"] and o["shape"][0] != 1:
                    problems.append(f"输出 batch 维 ≠ 1: {o['shape']}")

        return {"ok": not problems, "reason": "；".join(problems),
                "inputs": in_desc, "outputs": out_desc}
    except Exception as e:
        return {"ok": False, "reason": f"解析失败: {e}", "inputs": [], "outputs": []}


def main():
    ap = argparse.ArgumentParser(description="observer 训练任务脚本（Go --task-json 驱动）")
    ap.add_argument("--task-json", required=True, help="任务参数 JSON 文件路径")
    args = ap.parse_args()

    with open(args.task_json, encoding="utf-8") as f:
        task = json.load(f)

    global _LOG
    base = os.path.abspath(task["workdir"])

    def task_path(key):
        """任务参数里的路径全部相对 workdir，这里绝对化（chdir 失败也能定位）"""
        p = task[key]
        return p if os.path.isabs(p) else os.path.join(base, p)

    try:
        os.chdir(base)
    except Exception:
        pass  # 目录不可用则后续训练必然失败，错误路径仍尽力写 result.json

    log_file = task_path("log_file")
    result_file = task_path("result_file")
    global _LOG
    _LOG = None
    try:
        Path(log_file).parent.mkdir(parents=True, exist_ok=True)
        _LOG = open(log_file, "a", encoding="utf-8", buffering=1)
    except Exception:
        pass  # 日志不可写（如 workdir 缺失）不阻断错误路径，traceback 仍走 stderr

    imgsz = int(task.get("imgsz") or 704)
    epochs = int(task.get("epochs") or 150)
    batch = int(task.get("batch") or 16)
    device = task.get("device") or "0"
    data = os.path.join(task["yolo"], "dataset.yaml")
    names = read_names(task["yolo"])

    try:
        from ultralytics import YOLO
        model = YOLO("yolov8n.pt")
        register_callback(model)
        model.train(
            data=data, imgsz=imgsz, epochs=epochs,
            patience=PATIENCE, batch=batch, device=device, workers=4,
            # project 必须绝对路径，相对路径会被拼到默认 runs/detect 下造成双层嵌套
            project=task_path("project"), name="train",
            exist_ok=True, plots=True, multi_scale=MULTI_SCALE,
        )
        trainer = getattr(model, "trainer", None)
        if trainer is None:
            raise RuntimeError("训练完成但无法获取 trainer（save_dir 未知）")
        save_dir = Path(trainer.save_dir)

        best_tflite = find_best_tflite(save_dir)
        if best_tflite is None:
            # ultralytics ≥8.4 中 tflite 为废弃格式（走 onnx2tf 链路），litert 是官方替代，产物同为 best.tflite
            model.export(format="litert", imgsz=imgsz)
            best_tflite = find_best_tflite(save_dir)
        if best_tflite is None:
            raise RuntimeError("tflite 导出失败：weights 目录下未找到任何 .tflite 产物")

        # tflite 产物自检：shape 异常也写进 result.json（ok=false），Go 侧据此置失败
        tflite_check = check_tflite(str(best_tflite), imgsz)
        if not tflite_check["ok"]:
            print(f"[tflite-check] FAIL: {tflite_check['reason']}", file=sys.stderr)

        # 结果指标取 best 轮（发布的是 best.pt），缺省回退末轮
        metrics = getattr(trainer, "best_metrics", None) or trainer.metrics
        clean_metrics = {}
        for k, v in (metrics or {}).items():
            if isinstance(v, (int, float)) and math.isfinite(v):
                clean_metrics[k] = round(float(v), 5)

        # best_tflite 相对 workdir，Go 侧按此路径拉取（不再打包 zip，仅回传 tflite）
        write_result(result_file, {
            "metrics": clean_metrics,
            "names": names,
            "best_tflite": os.path.relpath(best_tflite, base),
            "tflite_check": tflite_check,
        })
    except Exception:
        err = traceback.format_exc()
        if _LOG is not None:
            _LOG.write("\n" + err + "\n")
        try:
            write_result(result_file, {"error": err[-2000:]})
        except Exception:
            pass
        sys.exit(1)


if __name__ == "__main__":
    main()