from __future__ import annotations

import argparse
from pathlib import Path

from ultralytics import YOLO


MODEL_ALIASES = {
    "YOLOv8n-seg": "yolov8n-seg.pt",
    "YOLOv8s-seg": "yolov8s-seg.pt",
    "YOLOv8m-seg": "yolov8m-seg.pt",
    "YOLOv8l-seg": "yolov8l-seg.pt",
    "YOLOv8x-seg": "yolov8x-seg.pt",
    "YOLOv9c-seg": "yolov9c-seg.pt",
    "YOLOv9e-seg": "yolov9e-seg.pt",
    "YOLO11n-seg": "yolo11n-seg.pt",
    "YOLO11s-seg": "yolo11s-seg.pt",
    "YOLO11m-seg": "yolo11m-seg.pt",
    "YOLO11l-seg": "yolo11l-seg.pt",
    "YOLO11x-seg": "yolo11x-seg.pt",
}


def resolve_model(value: str) -> str:
    candidate = MODEL_ALIASES.get(value, value)
    path = Path(candidate).expanduser()
    if path.is_absolute() or path.exists():
        return str(path.resolve())
    return candidate


def main() -> None:
    parser = argparse.ArgumentParser(description="Train a YOLO segmentation model from a supplied dataset.yaml.")
    parser.add_argument("--data", required=True, help="Path to YOLO dataset.yaml.")
    parser.add_argument("--model", default="YOLO11n-seg", help="Model alias, weight path, or Ultralytics model name.")
    parser.add_argument("--epochs", type=int, default=10)
    parser.add_argument("--imgsz", type=int, default=640)
    parser.add_argument("--batch", type=int, default=1)
    parser.add_argument("--workers", type=int, default=0)
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--project", required=True)
    parser.add_argument("--name", default="custom_upload")
    parser.add_argument("--exist-ok", action="store_true")
    args = parser.parse_args()

    model = YOLO(resolve_model(args.model))
    result = model.train(
        data=str(Path(args.data).expanduser().resolve()),
        epochs=args.epochs,
        imgsz=args.imgsz,
        batch=args.batch,
        workers=args.workers,
        device=args.device,
        project=str(Path(args.project).expanduser().resolve()),
        name=args.name,
        exist_ok=args.exist_ok,
        verbose=True,
    )
    save_dir = getattr(result, "save_dir", None) or getattr(model.trainer, "save_dir", "")
    print(f"save_dir={save_dir}")


if __name__ == "__main__":
    main()