# backend/ml/train_detector_model.py

import argparse
import json
import shutil
from pathlib import Path

from ultralytics import YOLO


def emit_progress(stage, progress, message="", **extra):
    out = {
        "type": "progress",
        "stage": stage,
        "progress": max(0.0, min(1.0, float(progress))),
        "message": message,
    }
    out.update(extra)
    print(json.dumps(out, ensure_ascii=False), flush=True)


def count_yolo_samples(dataset_root: Path, split: str) -> int:
    images_dir = dataset_root / "images" / split
    labels_dir = dataset_root / "labels" / split

    if not images_dir.exists() or not labels_dir.exists():
        return 0

    image_exts = {".jpg", ".jpeg", ".png", ".webp"}
    count = 0

    for image_path in images_dir.iterdir():
        if not image_path.is_file():
            continue

        if image_path.suffix.lower() not in image_exts:
            continue

        label_path = labels_dir / f"{image_path.stem}.txt"
        if label_path.exists() and label_path.stat().st_size > 0:
            count += 1

    return count


def safe_int(value, fallback):
    try:
        return int(value)
    except Exception:
        return fallback


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--root", required=True)
    parser.add_argument("--base", default="yolo26n.pt")
    parser.add_argument("--epochs", default="80")
    parser.add_argument("--imgsz", default="640")
    parser.add_argument("--device", default="auto")
    parser.add_argument("--workers", default="2")
    parser.add_argument("--patience", default="20")
    args = parser.parse_args()

    import torch

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

    epochs = max(1, safe_int(args.epochs, 80))
    imgsz = max(64, safe_int(args.imgsz, 640))
    workers = max(0, safe_int(args.workers, 2))
    patience = max(0, safe_int(args.patience, 20))

    if str(args.device).lower() == "auto":
        train_device = 0 if torch.cuda.is_available() else "cpu"
    else:
        train_device = args.device

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

    train_count = count_yolo_samples(dataset_root, "train")
    val_count = count_yolo_samples(dataset_root, "val")

    emit_progress(
        "detector",
        0.01,
        "YOLO-Dataset wird geprüft…",
        trainSamples=train_count,
        valSamples=val_count,
        epochs=epochs,
        imgsz=imgsz,
        device=str(train_device),
    )

    if train_count <= 0:
        raise SystemExit("no YOLO train samples found")

    if val_count <= 0:
        raise SystemExit("no YOLO val samples found")

    emit_progress(
        "detector",
        0.03,
        "YOLO-Basismodell wird geladen…",
        base=args.base,
        device=str(train_device),
    )

    model = YOLO(args.base)

    best_epoch = 0
    last_epoch = 0

    def on_train_epoch_start(trainer):
        epoch = int(getattr(trainer, "epoch", 0)) + 1
        total = int(getattr(trainer, "epochs", epochs) or epochs)

        emit_progress(
            "detector",
            0.04 + 0.90 * ((epoch - 1) / max(1, total)),
            f"Object Detector trainiert…",
            epoch=epoch,
            epochs=total,
            trainSamples=train_count,
            valSamples=val_count,
            device=str(train_device),
        )

    def on_train_epoch_end(trainer):
        nonlocal last_epoch

        epoch = int(getattr(trainer, "epoch", 0)) + 1
        total = int(getattr(trainer, "epochs", epochs) or epochs)
        last_epoch = max(last_epoch, epoch)

        emit_progress(
            "detector",
            0.04 + 0.90 * (epoch / max(1, total)),
            f"Object Detector trainiert…",
            epoch=epoch,
            epochs=total,
            trainSamples=train_count,
            valSamples=val_count,
            device=str(train_device),
        )

    def on_fit_epoch_end(trainer):
        nonlocal best_epoch

        epoch = int(getattr(trainer, "epoch", 0)) + 1
        total = int(getattr(trainer, "epochs", epochs) or epochs)
        metrics = getattr(trainer, "metrics", None) or {}

        # Ultralytics nutzt je nach Version unterschiedliche Keys.
        map50 = (
            metrics.get("metrics/mAP50(B)")
            or metrics.get("metrics/mAP50")
            or metrics.get("mAP50")
        )
        map5095 = (
            metrics.get("metrics/mAP50-95(B)")
            or metrics.get("metrics/mAP50-95")
            or metrics.get("mAP50-95")
        )

        if map50 is not None or map5095 is not None:
            best_epoch = epoch

        emit_progress(
            "detector",
            0.04 + 0.90 * (epoch / max(1, total)),
            f"Object Detector validiert…",
            epoch=epoch,
            epochs=total,
            mAP50=map50,
            mAP5095=map5095,
            device=str(train_device),
        )

    model.add_callback("on_train_epoch_start", on_train_epoch_start)
    model.add_callback("on_train_epoch_end", on_train_epoch_end)
    model.add_callback("on_fit_epoch_end", on_fit_epoch_end)

    emit_progress(
        "detector",
        0.05,
        "Object Detector Training startet…",
        trainSamples=train_count,
        valSamples=val_count,
        epochs=epochs,
        imgsz=imgsz,
        device=str(train_device),
    )

    result = model.train(
        data=str(yaml_path),
        epochs=epochs,
        imgsz=imgsz,
        project=str(runs_dir),
        name="detect",
        exist_ok=True,
        device=train_device,
        workers=workers,
        patience=patience,
    )

    emit_progress(
        "detector",
        0.96,
        "Bestes YOLO-Modell wird übernommen…",
        lastEpoch=last_epoch,
        bestEpoch=best_epoch,
        device=str(train_device),
    )

    best = runs_dir / "detect" / "weights" / "best.pt"
    last = runs_dir / "detect" / "weights" / "last.pt"

    if not best.exists():
        if last.exists():
            best = last
        else:
            raise SystemExit(f"best.pt not found after training: {best}")

    out_dir.mkdir(parents=True, exist_ok=True)

    final_model = out_dir / "best.pt"
    shutil.copy2(best, final_model)

    result_path = runs_dir / "detect"
    status = {
        "ok": True,
        "model": str(final_model),
        "sourceModel": str(best),
        "runs": str(result_path),
        "trainSamples": train_count,
        "valSamples": val_count,
        "epochs": epochs,
        "imgsz": imgsz,
        "device": str(train_device),
    }

    with (out_dir / "status.json").open("w", encoding="utf-8") as f:
        json.dump(status, f, ensure_ascii=False, indent=2)

    emit_progress(
        "detector",
        1.0,
        "Object Detector fertig.",
        **status,
    )

    print(json.dumps(status, ensure_ascii=False), flush=True)


if __name__ == "__main__":
    main()