# backend/ml/train_detector_model.py

import argparse
import json
import shutil
from datetime import datetime, timezone
from pathlib import Path

from ultralytics import YOLO
from ultralytics.models.yolo.detect.train import DetectionTrainer


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"
        # Eine vorhandene, leere Labeldatei ist ein gültiges YOLO-Negativbeispiel.
        if label_path.exists() and label_path.is_file():
            count += 1

    return count


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


def load_base_model(base):
    base_text = str(base or "").strip()
    base_path = Path(base_text).expanduser()

    is_explicit_path = base_path.is_absolute() or base_path.parent != Path(".")
    base_exists = base_path.is_file() and base_path.stat().st_size > 0
    if not is_explicit_path or base_exists:
        return YOLO(base_text), base_text

    base_path.parent.mkdir(parents=True, exist_ok=True)
    model_name = base_path.name
    model = YOLO(model_name)

    candidates = []
    ckpt_path = getattr(model, "ckpt_path", None)
    if ckpt_path:
        candidates.append(Path(str(ckpt_path)).expanduser())

    candidates.append(Path.cwd() / model_name)
    candidates.append(Path(model_name))

    copied = False
    for candidate in candidates:
        try:
            candidate = candidate.resolve()
            target = base_path.resolve()
        except Exception:
            continue

        if not candidate.exists() or candidate == target:
            continue

        try:
            shutil.copy2(candidate, target)
            copied = True

            if candidate.parent == Path.cwd().resolve():
                try:
                    candidate.unlink()
                except Exception:
                    pass
            break
        except Exception:
            continue

    if copied and base_path.exists():
        return YOLO(str(base_path)), str(base_path)

    return model, model_name


def batch_sample_ids(batch):
    if not isinstance(batch, dict):
        return []

    image_paths = batch.get("im_file")
    if not image_paths:
        return []

    if isinstance(image_paths, (str, Path)):
        image_paths = [image_paths]

    out = []
    seen = set()

    for image_path in image_paths:
        stem = Path(str(image_path)).stem.strip()
        if stem and stem not in seen:
            seen.add(stem)
            out.append(stem)

    return out


def progress_detection_trainer(train_count, val_count, train_device, fallback_epochs):
    class ProgressDetectionTrainer(DetectionTrainer):
        def __init__(self, *args, **kwargs):
            super().__init__(*args, **kwargs)
            self._preview_epoch = 0
            self._preview_batch = 0

        def preprocess_batch(self, batch):
            epoch = int(getattr(self, "epoch", 0)) + 1
            total_epochs = int(getattr(self, "epochs", fallback_epochs) or fallback_epochs)

            if epoch != self._preview_epoch:
                self._preview_epoch = epoch
                self._preview_batch = 0

            self._preview_batch += 1

            # Pro Batch das gerade trainierte Bild melden – ohne Drosselung, damit
            # das Frontend live immer das aktuell trainierte Bild anzeigen kann.
            sample_ids = batch_sample_ids(batch)
            if sample_ids:
                total_batches = max(1, len(self.train_loader))
                completed = (epoch - 1) + min(1.0, self._preview_batch / total_batches)
                progress = 0.04 + 0.90 * (completed / max(1, total_epochs))

                emit_progress(
                    "detector",
                    progress,
                    "Object Detector trainiert…",
                    epoch=epoch,
                    epochs=total_epochs,
                    sampleId=sample_ids[0],
                    trainSamples=train_count,
                    valSamples=val_count,
                    device=str(train_device),
                )

            return super().preprocess_batch(batch)

    return ProgressDetectionTrainer


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("--batch", default="0")
    parser.add_argument("--threads", default="0")
    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))
    batch_size = max(0, safe_int(args.batch, 0))
    threads = max(0, safe_int(args.threads, 0))
    patience = max(0, safe_int(args.patience, 20))

    if threads > 0:
        torch.set_num_threads(threads)
        try:
            torch.set_num_interop_threads(max(1, min(threads, 4)))
        except Exception:
            pass

    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, base_used = load_base_model(args.base)

    best_epoch = 0
    last_epoch = 0
    best_map50 = 0.0
    best_map5095 = 0.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, best_map50, best_map5095

        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

            # Bestes Modell merken (YOLO speichert best.pt nach Fitness ~ mAP).
            m50 = float(map50) if map50 is not None else 0.0
            if m50 >= best_map50:
                best_map50 = m50
                if map5095 is not None:
                    best_map5095 = float(map5095)

        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),
    )

    train_kwargs = {
        "trainer": progress_detection_trainer(
            train_count,
            val_count,
            train_device,
            epochs,
        ),
        "data": str(yaml_path),
        "epochs": epochs,
        "imgsz": imgsz,
        "project": str(runs_dir),
        "name": "detect",
        "exist_ok": True,
        "device": train_device,
        "workers": workers,
        "patience": patience,
    }
    if batch_size > 0:
        train_kwargs["batch"] = batch_size

    result = model.train(**train_kwargs)

    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),
        "baseModel": str(base_used),
        "runs": str(result_path),
        "trainedAt": datetime.now(timezone.utc).isoformat(),
        "trainSamples": train_count,
        "valSamples": val_count,
        "epochs": epochs,
        "imgsz": imgsz,
        "batchSize": batch_size,
        "patience": patience,
        "threads": threads,
        "workers": workers,
        "device": str(train_device),
        "mAP50": round(best_map50, 4),
        "mAP5095": round(best_map5095, 4),
        "bestEpoch": best_epoch,
    }

    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()