"""
Train YOLOv8s on Matt Nudi's bee/drone/queen dataset using Modal GPU.

The Modal container downloads the dataset itself via Roboflow (we pass
the API key from the local .env), trains for 50 epochs, and persists
the best.pt weights to a Modal Volume.

To run:
    py scripts/train_yolo_on_modal.py

After training, download the weights with:
    modal volume get apiarist-weights /apiarist/weights/best.pt weights/honey_bee_detector.pt
"""

import os
from pathlib import Path

import modal


APP_NAME = "apiarist-yolo-train"
VOLUME_NAME = "apiarist-weights"
EPOCHS = 60
IMG_SIZE = 640
BATCH = 32
BASE_WEIGHTS = "yolov8s.pt"

# Switched from matt-nudi (909 imgs, ~50 queens) to hendricks_ricky (3308 imgs,
# 892 queens, includes Varroa mite class). Should massively improve queen
# detection mAP and give us a real specialist for mites too.
DATASET_WORKSPACE = "hendricks_ricky-hotmail-de"
DATASET_PROJECT = "bee-project"
DATASET_VERSION = 2

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "ultralytics==8.3.81",
        "roboflow==1.1.50",
        "pyyaml",
    )
    .apt_install("libgl1", "libglib2.0-0")
)

vol = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)

app = modal.App(APP_NAME)


@app.function(
    image=image,
    gpu="T4",
    volumes={"/weights": vol},
    timeout=3 * 60 * 60,
)
def train(rf_api_key: str) -> str:
    import shutil
    import sys
    from pathlib import Path

    from roboflow import Roboflow
    from ultralytics import YOLO

    print("=" * 60)
    print(
        f"Downloading {DATASET_WORKSPACE}/{DATASET_PROJECT} "
        f"v{DATASET_VERSION} from Roboflow ..."
    )
    print("=" * 60)
    rf = Roboflow(api_key=rf_api_key)
    project = rf.workspace(DATASET_WORKSPACE).project(DATASET_PROJECT)
    version = project.version(DATASET_VERSION)
    dataset = version.download("yolov8", location="/tmp/dataset")
    print(f"Dataset ready at {dataset.location}")

    print("\n" + "=" * 60)
    print(f"Training YOLOv8s for {EPOCHS} epochs on T4 ...")
    print("=" * 60)
    model = YOLO(BASE_WEIGHTS)
    results = model.train(
        data=f"{dataset.location}/data.yaml",
        epochs=EPOCHS,
        imgsz=IMG_SIZE,
        batch=BATCH,
        project="/weights",
        name="apiarist",
        exist_ok=True,
        device=0,
        patience=15,
    )

    best_pt = Path("/weights/apiarist/weights/best.pt")
    if not best_pt.exists():
        print("ERROR: best.pt not found after training", file=sys.stderr)
        sys.exit(1)

    size_mb = best_pt.stat().st_size / 1024 / 1024
    print(f"\n[OK] best.pt saved at {best_pt} ({size_mb:.1f} MB)")

    vol.commit()
    return str(best_pt)


@app.local_entrypoint()
def main() -> None:
    from dotenv import load_dotenv

    load_dotenv()
    api_key = os.environ.get("ROBOFLOW_API_KEY")
    if not api_key:
        raise SystemExit(
            "Missing ROBOFLOW_API_KEY in .env. Add it before running."
        )

    print("Kicking off Modal training (this takes ~30-45 min on T4) ...")
    weights_path = train.remote(rf_api_key=api_key)
    print("\n" + "=" * 60)
    print(f"DONE. Weights at: {weights_path}")
    print("=" * 60)
    print(
        "\nDownload locally with:\n"
        f" modal volume get {VOLUME_NAME} /apiarist/weights/best.pt "
        f"weights/honey_bee_detector.pt"
    )