#!/usr/bin/env python3
"""CUDA-train an isolated, non-promotable YOLO candidate."""

from __future__ import annotations

import argparse
import hashlib
import json
import shutil
from pathlib import Path

EXPERIMENTAL_ROOT = Path("/app/storage/training/experimental")


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _inside(path: Path, root: Path) -> bool:
    try:
        path.resolve().relative_to(root.resolve())
        return True
    except ValueError:
        return False


def validate_paths(dataset_dir: Path, output_dir: Path) -> dict:
    if not _inside(dataset_dir, EXPERIMENTAL_ROOT) or not _inside(output_dir, EXPERIMENTAL_ROOT):
        raise ValueError(f"dataset and output must remain below {EXPERIMENTAL_ROOT}")
    marker_path = dataset_dir / "EXPERIMENTAL_ONLY.json"
    yaml_path = dataset_dir / "dataset.yaml"
    if not marker_path.is_file() or not yaml_path.is_file():
        raise ValueError("experimental marker or dataset.yaml is missing")
    marker = json.loads(marker_path.read_text(encoding="utf-8"))
    if marker.get("promotion_allowed") is not False or marker.get("release_claim_allowed") is not False:
        raise ValueError("experimental marker does not prohibit promotion and release claims")
    if marker.get("dataset_yaml_sha256") != sha256(yaml_path):
        raise ValueError("experimental dataset.yaml checksum mismatch")
    return marker


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--dataset-dir", type=Path, required=True)
    parser.add_argument("--base-model", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    parser.add_argument("--run-name", default="v73-flanders-remediation")
    parser.add_argument("--epochs", type=int, default=20)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--batch", type=int, default=8)
    parser.add_argument("--learning-rate", type=float, default=0.0001)
    parser.add_argument("--freeze", type=int, default=10)
    parser.add_argument("--workers", type=int, default=0)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    validate_paths(args.dataset_dir, args.output_dir)
    if not args.base_model.is_file():
        raise SystemExit(f"base model missing: {args.base_model}")
    import torch
    from ultralytics import YOLO

    if not torch.cuda.is_available():
        raise SystemExit("CUDA is required for GeoIntel experimental training")
    model = YOLO(str(args.base_model))
    result = model.train(
        data=str(args.dataset_dir / "dataset.yaml"),
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        workers=args.workers,
        device=0,
        project=str(args.output_dir),
        name=args.run_name,
        exist_ok=False,
        pretrained=True,
        deterministic=True,
        seed=73,
        amp=False,
        val=False,
        optimizer="AdamW",
        lr0=args.learning_rate,
        lrf=0.1,
        freeze=args.freeze,
        mosaic=0.0,
        translate=0.05,
        scale=0.1,
        plots=False,
        verbose=True,
    )
    run_dir = Path(result.save_dir)
    best = run_dir / "weights" / "best.pt"
    if not best.is_file():
        best = run_dir / "weights" / "last.pt"
    candidate = run_dir / "candidate.experimental.pt"
    shutil.copy2(best, candidate)
    summary = {
        "schema_version": 1,
        "status": "trained_experimental_only",
        "promotion_allowed": False,
        "release_claim_allowed": False,
        "human_review_pending": True,
        "dataset_yaml_sha256": sha256(args.dataset_dir / "dataset.yaml"),
        "base_model": str(args.base_model),
        "base_model_sha256": sha256(args.base_model),
        "candidate_model": str(candidate),
        "candidate_model_sha256": sha256(candidate),
        "epochs": args.epochs,
        "imgsz": args.imgsz,
        "batch": args.batch,
        "learning_rate": args.learning_rate,
        "optimizer": "AdamW",
        "freeze": args.freeze,
        "workers": args.workers,
        "device": torch.cuda.get_device_name(0),
        "required_next_gate": "independent frozen V72 calibration evaluation",
    }
    (run_dir / "experimental_training_summary.json").write_text(
        json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8"
    )
    print(json.dumps(summary, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())