#!/usr/bin/env python3
"""SageMaker entry point for YOLO electric meter detection training."""

from __future__ import annotations

import argparse
import json
import os
import shutil
from pathlib import Path
from typing import Any

import yaml
from sanitize_onnx import sanitize_onnx
from ultralytics import YOLO  # type: ignore[reportMissingImports]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="yolo26n.pt")
    parser.add_argument("--epochs", type=int, default=25)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--batch", type=int, default=16)
    parser.add_argument("--workers", type=int, default=2)
    parser.add_argument("--patience", type=int, default=20)
    parser.add_argument("--device", default=None)
    parser.add_argument("--data-yaml", default=None)
    parser.add_argument("--dataset-dir", default=os.environ.get("SM_CHANNEL_TRAIN", "/opt/ml/input/data/train"))
    parser.add_argument("--train-dir", dest="dataset_dir", help=argparse.SUPPRESS)
    parser.add_argument("--model-dir", default=os.environ.get("SM_MODEL_DIR", "/opt/ml/model"))
    return parser.parse_args()


def find_data_yaml(dataset_dir: Path, explicit_path: str | None) -> Path:
    if explicit_path:
        data_yaml = Path(explicit_path)
        if data_yaml.is_file():
            return data_yaml
        raise FileNotFoundError(f"Configured data.yaml does not exist: {data_yaml}")

    matches = sorted(dataset_dir.rglob("data.yaml"))
    if not matches:
        raise FileNotFoundError(f"Could not find data.yaml under {dataset_dir}")
    if len(matches) > 1:
        print(f"Found multiple data.yaml files; using {matches[0]}")
    return matches[0]


def prepare_data_yaml(data_yaml: Path) -> Path:
    """Write a SageMaker-local data file rooted at the uploaded dataset."""
    dataset_root = data_yaml.parent
    data = yaml.safe_load(data_yaml.read_text(encoding="utf-8"))
    if not isinstance(data, dict):
        raise ValueError(f"Expected a mapping in {data_yaml}")

    normalized = dict(data)
    normalized["path"] = str(dataset_root)
    if "val" not in normalized and "valid" in normalized:
        normalized["val"] = normalized.pop("valid")

    prepared_path = dataset_root / "data.sagemaker.yaml"
    prepared_path.write_text(yaml.safe_dump(normalized, sort_keys=False), encoding="utf-8")
    print(f"Prepared dataset config: {prepared_path}")
    return prepared_path


def copy_if_exists(source: Path, destination: Path) -> None:
    if source.exists():
        shutil.copy2(source, destination)
        print(f"Saved {destination}")


def main() -> None:
    args = parse_args()
    dataset_dir = Path(args.dataset_dir)
    model_dir = Path(args.model_dir)
    model_dir.mkdir(parents=True, exist_ok=True)

    data_yaml = prepare_data_yaml(find_data_yaml(dataset_dir, args.data_yaml))
    model = YOLO(args.model)

    train_kwargs: dict[str, Any] = {
        "data": str(data_yaml),
        "epochs": args.epochs,
        "imgsz": args.imgsz,
        "batch": args.batch,
        "workers": args.workers,
        "patience": args.patience,
        "project": str(model_dir / "runs"),
        "name": "train",
        "exist_ok": True,
    }
    if args.device:
        train_kwargs["device"] = args.device

    results = model.train(**train_kwargs)
    save_dir = Path(results.save_dir)
    best_pt = save_dir / "weights" / "best.pt"
    last_pt = save_dir / "weights" / "last.pt"
    trained_weights = best_pt if best_pt.exists() else last_pt
    if not trained_weights.exists():
        raise FileNotFoundError(f"Could not find trained weights in {save_dir / 'weights'}")

    copy_if_exists(trained_weights, model_dir / "best.pt")
    trained_model = YOLO(str(trained_weights))
    onnx_path = Path(trained_model.export(format="onnx", imgsz=args.imgsz))
    saved_onnx_path = sanitize_onnx(onnx_path, model_dir / "model.onnx")
    print(f"Saved {saved_onnx_path}")

    metrics = {
        "model": args.model,
        "epochs": args.epochs,
        "imgsz": args.imgsz,
        "batch": args.batch,
        "workers": args.workers,
        "patience": args.patience,
        "data_yaml": str(data_yaml),
        "weights": str(trained_weights),
        "onnx": str(saved_onnx_path),
    }
    (model_dir / "metrics.json").write_text(json.dumps(metrics, indent=2), encoding="utf-8")
    print(f"Saved model artifacts to {model_dir}")


if __name__ == "__main__":
    main()