# backend\ml\train_detector_model.py

import argparse
import json
from pathlib import Path
from ultralytics import YOLO


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--base", default="yolo11n.pt")
    parser.add_argument("--epochs", default="80")
    parser.add_argument("--imgsz", default="640")
    args = parser.parse_args()

    root = Path(args.root).resolve()
    yaml_path = root / "detector" / "dataset" / "dataset.yaml"
    runs_dir = root / "detector" / "runs"

    if not yaml_path.exists():
        raise SystemExit(f"dataset.yaml not found: {yaml_path}")

    model = YOLO(args.base)

    result = model.train(
        data=str(yaml_path),
        epochs=int(args.epochs),
        imgsz=int(args.imgsz),
        project=str(runs_dir),
        name="detect",
        exist_ok=True,
        device="cpu",
        workers=2,
        patience=20,
    )

    best = runs_dir / "detect" / "weights" / "best.pt"
    if not best.exists():
        raise SystemExit(f"best.pt not found after training: {best}")

    out_dir = root / "detector" / "model"
    out_dir.mkdir(parents=True, exist_ok=True)

    final_model = out_dir / "best.pt"
    final_model.write_bytes(best.read_bytes())

    print(json.dumps({
        "ok": True,
        "model": str(final_model),
        "runs": str(runs_dir / "detect"),
    }))


if __name__ == "__main__":
    main()