#!/usr/bin/env python3
"""Cloud GPU training script for YOLO12n 2-class checkbox detector.

Single-file, idempotent script with three stages:
  1. prepare_data — Extract tarball, tile train/val into overlapping crops
  2. train        — Train YOLO12n with auto-detected device and cache strategy
  3. export       — Export best.pt to ONNX with shape validation

All config via CLI args. Designed to be scp'd to a RunPod A100 and run via SSH.

Usage:
    # Full pipeline
    python train_cloud.py --tarball /workspace/yolo_2class_slim.tar.gz

    # Smoke test (3 epochs, small batch)
    python train_cloud.py --tarball /workspace/yolo_2class_slim.tar.gz --epochs 3 --batch 8

    # Stage control
    python train_cloud.py --tarball ... --tile-only
    python train_cloud.py --data-dir /workspace/data/yolo_2class_slim --train-only
    python train_cloud.py --tarball ... --export-only /path/to/best.pt

    # Override defaults
    python train_cloud.py --tarball ... --tile-size 1216 --imgsz 1216 --batch 48
    python train_cloud.py --tarball ... --base-model yolo12s.pt --batch 24
    python train_cloud.py --tarball ... --seed 123
"""

import argparse
import hashlib
import json
import logging
import os
import platform
import random
import shutil
import subprocess
import sys
import tarfile
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple

from PIL import Image
from tqdm import tqdm

logger = logging.getLogger(__name__)

# ── Defaults ──────────────────────────────────────────────────────────────────

VOLUME_DIR = "/runpod-volume"
DEFAULT_WORK_DIR = "/workspace/data"
DEFAULT_OUTPUT_DIR = f"{VOLUME_DIR}/output"
DEFAULT_TILE_SIZE = 1024
DEFAULT_TILE_OVERLAP = 0.2
DEFAULT_MIN_BOX_VISIBLE = 0.4
DEFAULT_MIN_BOX_PX = 4
DEFAULT_KEEP_EMPTY_RATIO = 0.15
DEFAULT_SEED = 42
DEFAULT_NUM_WORKERS = min(8, os.cpu_count() or 4)

DEFAULT_BASE_MODEL = "yolo12n.pt"
DEFAULT_IMGSZ = 1024
DEFAULT_EPOCHS = 200
DEFAULT_BATCH = 48
DEFAULT_PATIENCE = 50
DEFAULT_CLS_LOSS = 1.5
SAVE_PERIOD = 25
ONNX_OPSET = 18

YOLO_CLASS_NAMES = ["checkbox_checked", "checkbox_unchecked"]

IMG_EXT = ".jpg"
IMG_QUALITY = 95


# ═══════════════════════════════════════════════════════════════════════════════
# Stage 1: Data Preparation
# ═══════════════════════════════════════════════════════════════════════════════


def _sha256(path: Path) -> str:
    """Compute SHA256 of a file."""
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def _generate_tiles(
    img_w: int,
    img_h: int,
    tile_size: int,
    overlap: float,
) -> List[Tuple[int, int, int, int]]:
    """Generate overlapping tile grid for an image.

    Returns list of (x0, y0, x1, y1) tile coordinates.
    Images smaller than tile_size in both dimensions return a single tile.
    """
    if img_w <= tile_size and img_h <= tile_size:
        return [(0, 0, img_w, img_h)]

    stride = max(1, int(tile_size * (1 - overlap)))
    tiles = []
    for y in range(0, img_h, stride):
        for x in range(0, img_w, stride):
            x1 = min(x + tile_size, img_w)
            y1 = min(y + tile_size, img_h)
            x0 = max(0, x1 - tile_size)
            y0 = max(0, y1 - tile_size)
            tiles.append((x0, y0, x1, y1))

    # Deduplicate (edge tiles can overlap with snapped-back tiles)
    return list(dict.fromkeys(tiles))


def _clip_annots_to_tile(
    labels: List[str],
    tile: Tuple[int, int, int, int],
    img_w: int,
    img_h: int,
    min_box_visible: float,
    min_box_px: int,
) -> List[str]:
    """Remap YOLO labels to tile-local coords. Returns new label lines."""
    tx0, ty0, tx1, ty1 = tile
    tile_w, tile_h = tx1 - tx0, ty1 - ty0
    out = []
    for line in labels:
        parts = line.strip().split()
        if len(parts) < 5:
            continue
        cls = parts[0]
        cx, cy, nw, nh = (
            float(parts[1]),
            float(parts[2]),
            float(parts[3]),
            float(parts[4]),
        )
        # Convert YOLO normalized to pixel coords
        ax = cx * img_w - (nw * img_w) / 2
        ay = cy * img_h - (nh * img_h) / 2
        aw = nw * img_w
        ah = nh * img_h

        # Intersection with tile
        ix0 = max(ax, tx0)
        iy0 = max(ay, ty0)
        ix1 = min(ax + aw, tx1)
        iy1 = min(ay + ah, ty1)

        if ix0 >= ix1 or iy0 >= iy1:
            continue

        # Check visibility fraction
        box_area = max(aw * ah, 1e-6)
        if (ix1 - ix0) * (iy1 - iy0) / box_area < min_box_visible:
            continue

        lw = ix1 - ix0
        lh = iy1 - iy0
        if lw < min_box_px or lh < min_box_px:
            continue

        # Convert back to YOLO normalized (tile-local)
        lcx = ((ix0 - tx0) + lw / 2) / tile_w
        lcy = ((iy0 - ty0) + lh / 2) / tile_h
        lnw = lw / tile_w
        lnh = lh / tile_h
        out.append(f"{cls} {lcx:.6f} {lcy:.6f} {lnw:.6f} {lnh:.6f}")
    return out


def _read_labels_safe(lbl_path: Path) -> List[str]:
    """Read label file, skipping corrupt/non-UTF-8 files."""
    if not lbl_path.exists():
        return []
    try:
        text = lbl_path.read_text(encoding="utf-8").strip()
        return text.split("\n") if text else []
    except (UnicodeDecodeError, ValueError):
        return []


def _process_single_image(args: tuple) -> Tuple[int, int, int]:
    """Tile a single image and write tiles + labels.

    Returns (tiles_written, empty_kept, empty_skipped).
    """
    (
        img_path,
        lbl_dir,
        out_img_dir,
        out_lbl_dir,
        rng_seed,
        tile_size,
        tile_overlap,
        min_box_visible,
        min_box_px,
        keep_empty_ratio,
    ) = args
    rng = random.Random(rng_seed)
    stem = img_path.stem

    # Skip macOS resource fork files
    if stem.startswith("._"):
        return 0, 0, 0

    lbl_path = lbl_dir / f"{stem}.txt"
    labels = _read_labels_safe(lbl_path)

    tiles_written = 0
    empty_kept = 0
    empty_skipped = 0

    with Image.open(img_path) as img:
        w, h = img.size
        grey = img.convert("L").convert("RGB")
        tiles = _generate_tiles(w, h, tile_size, tile_overlap)

        for ti, tile in enumerate(tiles):
            tile_labels = _clip_annots_to_tile(
                labels, tile, w, h, min_box_visible, min_box_px
            )
            tile_name = f"{stem}_t{ti}"

            if not tile_labels:
                if rng.random() < keep_empty_ratio:
                    tx0, ty0, tx1, ty1 = tile
                    tile_img = grey.crop((tx0, ty0, tx1, ty1))
                    tile_img.save(
                        out_img_dir / f"{tile_name}{IMG_EXT}", quality=IMG_QUALITY
                    )
                    (out_lbl_dir / f"{tile_name}.txt").write_text("")
                    empty_kept += 1
                    tiles_written += 1
                else:
                    empty_skipped += 1
                continue

            tx0, ty0, tx1, ty1 = tile
            tile_img = grey.crop((tx0, ty0, tx1, ty1))
            tile_img.save(
                out_img_dir / f"{tile_name}{IMG_EXT}", quality=IMG_QUALITY
            )
            (out_lbl_dir / f"{tile_name}.txt").write_text(
                "\n".join(tile_labels) + "\n"
            )
            tiles_written += 1

    return tiles_written, empty_kept, empty_skipped


def _tile_split(
    split_dir: Path,
    tile_size: int,
    tile_overlap: float,
    min_box_visible: float,
    min_box_px: int,
    keep_empty_ratio: float,
    seed: int,
    num_workers: int,
) -> Dict[str, int]:
    """Tile all images in a split directory. Atomic swap: tiles replace originals."""
    img_dir = split_dir / "images"
    lbl_dir = split_dir / "labels"
    split_name = split_dir.name

    # Guard: skip if already tiled
    existing = [f for f in img_dir.iterdir() if f.suffix in (".png", ".jpg")]
    if any(f.stem.endswith("_t0") for f in existing):
        logger.info("%s: already tiled, skipping", split_name)
        return {"skipped": True}

    # Detect source extension
    png_count = sum(1 for f in existing if f.suffix == ".png")
    jpg_count = sum(1 for f in existing if f.suffix == ".jpg")
    src_ext = ".png" if png_count >= jpg_count else ".jpg"

    out_base = split_dir.parent / f"{split_name}_tiled"
    if out_base.exists():
        shutil.rmtree(out_base)  # Clean stale data from prior crashed run
    out_img_dir = out_base / "images"
    out_lbl_dir = out_base / "labels"
    out_img_dir.mkdir(parents=True, exist_ok=True)
    out_lbl_dir.mkdir(parents=True, exist_ok=True)

    images = sorted(img_dir.glob(f"*{src_ext}"))
    if not images:
        # Try the other extension
        other_ext = ".jpg" if src_ext == ".png" else ".png"
        images = sorted(img_dir.glob(f"*{other_ext}"))

    base_rng = random.Random(seed)
    work = [
        (
            img_path,
            lbl_dir,
            out_img_dir,
            out_lbl_dir,
            base_rng.randint(0, 2**31),
            tile_size,
            tile_overlap,
            min_box_visible,
            min_box_px,
            keep_empty_ratio,
        )
        for img_path in images
    ]

    total_tiles = 0
    kept_empty = 0
    skipped_empty = 0

    with ThreadPoolExecutor(max_workers=num_workers) as pool:
        futures = {pool.submit(_process_single_image, w): w for w in work}
        for future in tqdm(
            as_completed(futures), total=len(futures), desc=f"Tiling {split_name}"
        ):
            tw, ek, es = future.result()
            total_tiles += tw
            kept_empty += ek
            skipped_empty += es

    # Atomic swap: remove original, rename tiled to original
    shutil.rmtree(split_dir)
    out_img_dir.parent.rename(split_dir)

    logger.info(
        "%s: %d pages -> %d tiles (%d empty/negative, %d dropped)",
        split_name,
        len(images),
        total_tiles,
        kept_empty,
        skipped_empty,
    )

    return {
        "source_images": len(images),
        "tiles": total_tiles,
        "empty_kept": kept_empty,
        "empty_skipped": skipped_empty,
    }


def _fix_dataset_yaml(data_dir: Path):
    """Fix dataset.yaml path to absolute (Ultralytics resolves relative to cwd)."""
    yaml_path = data_dir / "dataset.yaml"
    if not yaml_path.exists():
        logger.warning("No dataset.yaml found at %s", yaml_path)
        return

    text = yaml_path.read_text()
    if "path: ." in text:
        text = text.replace("path: .", f"path: {data_dir}")
        yaml_path.write_text(text)
        logger.info("Fixed dataset.yaml path to %s", data_dir)


def _count_images(split_dir: Path) -> int:
    """Count images in a split directory."""
    img_dir = split_dir / "images"
    if not img_dir.exists():
        return 0
    return sum(1 for f in img_dir.iterdir() if f.suffix in (".png", ".jpg"))


def _get_environment_info() -> Dict[str, Any]:
    """Collect environment info for the manifest."""
    info: Dict[str, Any] = {
        "platform": platform.platform(),
        "python_version": sys.version,
        "hostname": platform.node(),
    }

    # Torch version
    try:
        import torch

        info["torch_version"] = torch.__version__
        info["cuda_available"] = torch.cuda.is_available()
        if torch.cuda.is_available():
            info["cuda_version"] = torch.version.cuda
            info["gpu_name"] = torch.cuda.get_device_name(0)
            info["gpu_vram_gb"] = round(
                torch.cuda.get_device_properties(0).total_memory / 1e9, 1
            )
    except ImportError:
        info["torch_version"] = "not installed"

    # Git commit
    try:
        result = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            capture_output=True,
            text=True,
            timeout=5,
        )
        if result.returncode == 0:
            info["git_commit"] = result.stdout.strip()
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    # pip freeze
    try:
        result = subprocess.run(
            [sys.executable, "-m", "pip", "freeze"],
            capture_output=True,
            text=True,
            timeout=30,
        )
        if result.returncode == 0:
            info["pip_freeze"] = result.stdout.strip().split("\n")
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass

    return info


def write_manifest(
    manifest_path: Path,
    args: argparse.Namespace,
    tarball_sha256: Optional[str],
    tile_stats: Dict[str, Any],
    data_dir: Path,
) -> Dict[str, Any]:
    """Write training_manifest.json with full reproducibility info."""
    manifest: Dict[str, Any] = {
        "created_at": datetime.utcnow().isoformat() + "Z",
        "tarball_sha256": tarball_sha256,
        "tile_params": {
            "tile_size": args.tile_size,
            "tile_overlap": args.tile_overlap,
            "min_box_visible": DEFAULT_MIN_BOX_VISIBLE,
            "min_box_px": DEFAULT_MIN_BOX_PX,
            "keep_empty_ratio": DEFAULT_KEEP_EMPTY_RATIO,
        },
        "seed": args.seed,
        "cli_args": vars(args),
        "tile_stats": tile_stats,
        "image_counts": {
            "train": _count_images(data_dir / "train"),
            "val": _count_images(data_dir / "val"),
            "test": _count_images(data_dir / "test"),
        },
        "environment": _get_environment_info(),
    }

    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(json.dumps(manifest, indent=2, default=str))
    logger.info("Wrote manifest: %s", manifest_path)
    return manifest


def _safe_extract(tar: tarfile.TarFile, dest: Path) -> None:
    """Extract tarball with path traversal and symlink protection."""
    dest_resolved = dest.resolve()
    members = tar.getmembers()
    for member in tqdm(members, desc="Extracting", unit="files"):
        if member.issym() or member.islnk():
            logger.warning("Skipping link in tarball: %s", member.name)
            continue
        target = (dest / member.name).resolve()
        if not str(target).startswith(str(dest_resolved) + os.sep) and target != dest_resolved:
            logger.warning("Skipping path traversal in tarball: %s", member.name)
            continue
        tar.extract(member, str(dest))


def prepare_data(args: argparse.Namespace) -> Path:
    """Stage 1: Extract tarball, tile train/val, fix dataset.yaml, write manifest.

    Returns the data directory path.
    """
    work_dir = Path(args.work_dir)
    work_dir.mkdir(parents=True, exist_ok=True)

    tarball = Path(args.tarball)
    if not tarball.exists():
        raise FileNotFoundError(f"Tarball not found: {tarball}")

    # Determine data dir from tarball name (e.g. yolo_2class_slim.tar.gz -> yolo_2class_slim)
    tar_stem = tarball.name
    for suffix in (".tar.gz", ".tgz", ".tar"):
        if tar_stem.endswith(suffix):
            tar_stem = tar_stem[: -len(suffix)]
            break
    data_dir = work_dir / tar_stem

    # ── Extract ───────────────────────────────────────────────────────────
    if data_dir.exists() and (data_dir / "dataset.yaml").exists():
        logger.info("Dataset already extracted at %s, skipping extraction", data_dir)
    else:
        logger.info("Extracting %s to %s ...", tarball, work_dir)
        with tarfile.open(str(tarball), "r:*") as tar:
            _safe_extract(tar, work_dir)
        logger.info("Extraction complete: %s", data_dir)

    # Clean macOS resource forks
    for f in data_dir.rglob("._*"):
        f.unlink()

    # ── Compute tarball hash ──────────────────────────────────────────────
    logger.info("Computing tarball SHA256 ...")
    tarball_sha256 = _sha256(tarball)
    logger.info("Tarball SHA256: %s", tarball_sha256)

    # ── Tile train/val ────────────────────────────────────────────────────
    tile_stats = {}
    for split in ("train", "val"):
        split_dir = data_dir / split
        if not split_dir.exists():
            logger.warning("Split directory not found: %s", split_dir)
            continue

        stats = _tile_split(
            split_dir=split_dir,
            tile_size=args.tile_size,
            tile_overlap=args.tile_overlap,
            min_box_visible=DEFAULT_MIN_BOX_VISIBLE,
            min_box_px=DEFAULT_MIN_BOX_PX,
            keep_empty_ratio=DEFAULT_KEEP_EMPTY_RATIO,
            seed=args.seed,
            num_workers=DEFAULT_NUM_WORKERS,
        )
        tile_stats[split] = stats

    # Test split is NEVER tiled
    test_dir = data_dir / "test"
    if test_dir.exists():
        tile_stats["test"] = {"source_images": _count_images(test_dir), "tiled": False}

    # ── Fix dataset.yaml ──────────────────────────────────────────────────
    _fix_dataset_yaml(data_dir)

    # ── Write manifest ────────────────────────────────────────────────────
    manifest_path = data_dir / "training_manifest.json"
    write_manifest(manifest_path, args, tarball_sha256, tile_stats, data_dir)

    return data_dir


# ═══════════════════════════════════════════════════════════════════════════════
# Stage 2: Training
# ═══════════════════════════════════════════════════════════════════════════════


def _detect_device() -> str:
    """Auto-detect best available training device: cuda -> mps -> cpu."""
    try:
        import torch

        if torch.cuda.is_available():
            logger.info("Detected CUDA device")
            return "0"
        if torch.backends.mps.is_available():
            logger.info("Detected MPS (Apple Silicon) device")
            return "mps"
    except Exception as e:
        logger.debug("Device auto-detection failed, falling back to cpu: %r", e)
    return "cpu"


def _detect_cache_strategy() -> str:
    """Auto-detect cache strategy based on available RAM."""
    try:
        import psutil

        total_gb = psutil.virtual_memory().total / (1024**3)
        if total_gb >= 100:
            logger.info("System RAM: %.0f GB — using cache='ram'", total_gb)
            return "ram"
        else:
            logger.info("System RAM: %.0f GB — using cache='disk'", total_gb)
            return "disk"
    except ImportError:
        logger.info("psutil not installed — using cache='disk'")
        return "disk"


def _find_resume_weights(runs_dir: Path) -> Optional[str]:
    """Find the last.pt checkpoint for resuming interrupted training."""
    if not runs_dir.exists():
        return None

    candidates = sorted(
        runs_dir.glob("**/weights/last.pt"),
        key=lambda p: p.stat().st_mtime,
        reverse=True,
    )
    return str(candidates[0]) if candidates else None


def _resolve_train_dir(model, results, resume_weights: Optional[str]) -> Path:
    """Reliably determine the training output directory."""
    save_dir = getattr(results, "save_dir", None)
    if save_dir:
        return Path(save_dir)

    trainer = getattr(model, "trainer", None)
    save_dir = getattr(trainer, "save_dir", None)
    if save_dir:
        return Path(save_dir)

    if resume_weights:
        return Path(resume_weights).resolve().parents[1]

    raise RuntimeError(
        "Could not determine training save_dir. "
        "Check Ultralytics version compatibility."
    )


def _log_metrics(results_csv: Path) -> Dict[str, float]:
    """Read final metrics from results.csv."""
    import csv

    metrics: Dict[str, float] = {}
    with open(results_csv, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        last = None
        for row in reader:
            last = row

    if not last:
        return metrics

    for key in last:
        clean = key.strip()
        if "mAP50" in clean and "mAP50-95" not in clean:
            try:
                metrics["mAP50"] = float(last[key])
            except (ValueError, TypeError):
                pass
        elif "mAP50-95" in clean:
            try:
                metrics["mAP50_95"] = float(last[key])
            except (ValueError, TypeError):
                pass

    return metrics


def train(args: argparse.Namespace, data_dir: Path) -> Dict[str, Any]:
    """Stage 2: Train YOLO12n checkbox detector.

    Returns dict with training stats and output paths.
    """
    try:
        from ultralytics import YOLO
    except ImportError:
        raise ImportError(
            "ultralytics is required for training. "
            "Install with: pip install ultralytics"
        )

    dataset_yaml = data_dir / "dataset.yaml"
    if not dataset_yaml.exists():
        raise FileNotFoundError(
            f"Dataset YAML not found: {dataset_yaml}\n"
            "Run prepare_data first (or use --tarball)."
        )

    # ── WandB (optional) ──────────────────────────────────────────────────
    try:
        import wandb  # noqa: F401
        from ultralytics import settings as ul_settings

        if os.getenv("WANDB_API_KEY"):
            wandb.login()
            ul_settings.update({"wandb": True})
            logger.info("WandB enabled — metrics will be logged")
        else:
            logger.info("WANDB_API_KEY not set — skipping WandB")
    except ImportError:
        logger.info("WandB not installed — skipping metric logging")
    except Exception as e:
        logger.warning("WandB setup failed (continuing without): %s", e)

    stats: Dict[str, Any] = {}
    device = _detect_device()
    cache = _detect_cache_strategy()
    resume_weights: Optional[str] = None

    # Per-run output directory
    run_name = f"run_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
    runs_dir = data_dir / "runs"

    # ── Resume or fresh start ─────────────────────────────────────────────
    if args.resume:
        resume_weights = _find_resume_weights(runs_dir)
        if resume_weights is None:
            logger.warning("No checkpoint found to resume from. Starting fresh.")
        else:
            logger.info("Resuming from checkpoint: %s", resume_weights)

    if resume_weights:
        model = YOLO(resume_weights)
        logger.info("Resuming training on device=%s ...", device)
        results = model.train(resume=True, device=device)
    else:
        model = YOLO(args.base_model)
        logger.info(
            "Starting training: model=%s, device=%s, imgsz=%d, batch=%d, epochs=%d, seed=%d",
            args.base_model,
            device,
            args.imgsz,
            args.batch,
            args.epochs,
            args.seed,
        )
        results = model.train(
            data=str(dataset_yaml),
            epochs=args.epochs,
            imgsz=args.imgsz,
            batch=args.batch,
            device=device,
            patience=args.patience,
            cos_lr=True,
            cls=DEFAULT_CLS_LOSS,
            save_period=SAVE_PERIOD,
            project=str(runs_dir),
            name=run_name,
            exist_ok=True,
            verbose=True,
            mosaic=0.0,  # Disabled: defeats tiling strategy
            scale=0.0,  # Disabled: random scaling would shrink checkboxes
            fliplr=0.0,  # Disabled: forms/documents are never mirrored
            flipud=0.0,  # Disabled: forms/documents are never upside-down
            deterministic=True,
            seed=args.seed,
            cache=cache,
            workers=DEFAULT_NUM_WORKERS,
        )

    # ── Find best weights ─────────────────────────────────────────────────
    train_dir = _resolve_train_dir(model, results, resume_weights)
    best_pt = train_dir / "weights" / "best.pt"
    if not best_pt.exists():
        raise FileNotFoundError(
            f"Training did not produce best.pt at {best_pt}. "
            "Check training logs above."
        )

    stats["best_pt"] = str(best_pt)
    stats["train_dir"] = str(train_dir)
    logger.info("Training complete. Best weights: %s", best_pt)

    # ── Read metrics ──────────────────────────────────────────────────────
    results_csv = train_dir / "results.csv"
    if results_csv.exists():
        metrics = _log_metrics(results_csv)
        stats.update(metrics)
        if "mAP50" in metrics:
            logger.info("Final mAP50: %.4f", metrics["mAP50"])
        if "mAP50_95" in metrics:
            logger.info("Final mAP50-95: %.4f", metrics["mAP50_95"])

    return stats


# ═══════════════════════════════════════════════════════════════════════════════
# Stage 3: Export
# ═══════════════════════════════════════════════════════════════════════════════


def _validate_onnx(onnx_path: str):
    """Sanity-check the exported ONNX model's input/output shapes.

    Lenient: warns on unexpected shapes but doesn't raise.
    """
    try:
        import onnxruntime as ort
    except ImportError:
        logger.warning("onnxruntime not installed — skipping ONNX validation")
        return

    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    inp = sess.get_inputs()[0]
    out = sess.get_outputs()[0]
    logger.info("ONNX input:  %s %s", inp.name, inp.shape)
    logger.info("ONNX output: %s %s", out.name, out.shape)

    if len(out.shape) != 3:
        logger.warning(
            "Unexpected ONNX output rank %d: %s; skipping strict validation",
            len(out.shape),
            out.shape,
        )
        return

    expected = 4 + len(YOLO_CLASS_NAMES)
    dim1 = out.shape[1]
    dim2 = out.shape[2]

    if isinstance(dim1, int) and dim1 == expected:
        num_classes = dim1 - 4
    elif isinstance(dim2, int) and dim2 == expected:
        num_classes = dim2 - 4
    else:
        logger.warning(
            "Unexpected ONNX output shape %s; skipping strict class-count check",
            out.shape,
        )
        return

    logger.info("ONNX validation passed (%d classes)", num_classes)


def export(args: argparse.Namespace, best_pt: str) -> str:
    """Stage 3: Export best.pt to ONNX and copy to output dir.

    Returns path to the final ONNX file.
    """
    from ultralytics import YOLO

    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    logger.info("Exporting %s to ONNX (imgsz=%d) ...", best_pt, args.imgsz)
    model = YOLO(best_pt)
    export_path = model.export(
        format="onnx",
        imgsz=args.imgsz,
        simplify=True,
        opset=ONNX_OPSET,
    )

    if isinstance(export_path, (list, tuple)):
        export_path = export_path[0]
    if not export_path or not Path(export_path).exists():
        raise FileNotFoundError(
            f"ONNX export failed or returned invalid path: {export_path}"
        )

    logger.info("ONNX exported: %s", export_path)
    _validate_onnx(str(export_path))

    # Copy to output dir
    onnx_dst = output_dir / "checkbox_yolo12n_2class.onnx"
    shutil.copy2(str(export_path), str(onnx_dst))
    logger.info("ONNX copied to: %s", onnx_dst)

    pt_dst = output_dir / "best.pt"
    shutil.copy2(best_pt, str(pt_dst))
    logger.info("best.pt copied to: %s", pt_dst)

    # Copy manifest if it exists alongside the data (walk up from best.pt)
    best_pt_path = Path(best_pt)
    for parent in best_pt_path.parents:
        manifest_src = parent / "training_manifest.json"
        if manifest_src.exists():
            shutil.copy2(str(manifest_src), str(output_dir / "training_manifest.json"))
            logger.info("Manifest copied to: %s", output_dir / "training_manifest.json")
            break

    return str(onnx_dst)


# ═══════════════════════════════════════════════════════════════════════════════
# CLI
# ═══════════════════════════════════════════════════════════════════════════════


def main():
    parser = argparse.ArgumentParser(
        description="Cloud GPU training for YOLO12n 2-class checkbox detector",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Examples:
  # Full pipeline
  python train_cloud.py --tarball /workspace/yolo_2class_slim.tar.gz

  # Smoke test
  python train_cloud.py --tarball /workspace/yolo_2class_slim.tar.gz --epochs 3 --batch 8

  # Tile only (verify data)
  python train_cloud.py --tarball /workspace/yolo_2class_slim.tar.gz --tile-only

  # Train only (data already tiled)
  python train_cloud.py --data-dir /workspace/data/yolo_2class_slim --train-only

  # Export only
  python train_cloud.py --tarball ... --export-only /path/to/best.pt
""",
    )

    # Data args
    parser.add_argument(
        "--tarball",
        type=str,
        help="Path to yolo_2class_slim.tar.gz",
    )
    parser.add_argument(
        "--data-dir",
        type=str,
        help="Path to already-extracted dataset (skips extraction)",
    )
    parser.add_argument(
        "--work-dir",
        type=str,
        default=DEFAULT_WORK_DIR,
        help=f"Working directory for extraction and tiling (default: {DEFAULT_WORK_DIR})",
    )
    parser.add_argument(
        "--output-dir",
        type=str,
        default=DEFAULT_OUTPUT_DIR,
        help=f"Directory for final ONNX + weights (default: {DEFAULT_OUTPUT_DIR})",
    )

    # Tiling args
    parser.add_argument(
        "--tile-size",
        type=int,
        default=DEFAULT_TILE_SIZE,
        help=f"Tile size in pixels (default: {DEFAULT_TILE_SIZE})",
    )
    parser.add_argument(
        "--tile-overlap",
        type=float,
        default=DEFAULT_TILE_OVERLAP,
        help=f"Tile overlap fraction (default: {DEFAULT_TILE_OVERLAP})",
    )
    parser.add_argument(
        "--seed",
        type=int,
        default=DEFAULT_SEED,
        help=f"Random seed for tiling RNG and training (default: {DEFAULT_SEED})",
    )

    # Training args
    parser.add_argument(
        "--epochs",
        type=int,
        default=DEFAULT_EPOCHS,
        help=f"Max training epochs (default: {DEFAULT_EPOCHS})",
    )
    parser.add_argument(
        "--batch",
        type=int,
        default=DEFAULT_BATCH,
        help=f"Batch size (default: {DEFAULT_BATCH})",
    )
    parser.add_argument(
        "--imgsz",
        type=int,
        default=DEFAULT_IMGSZ,
        help=f"Model input image size (default: {DEFAULT_IMGSZ})",
    )
    parser.add_argument(
        "--patience",
        type=int,
        default=DEFAULT_PATIENCE,
        help=f"Early stopping patience (default: {DEFAULT_PATIENCE})",
    )
    parser.add_argument(
        "--base-model",
        type=str,
        default=DEFAULT_BASE_MODEL,
        help=f"Pretrained model to start from (default: {DEFAULT_BASE_MODEL})",
    )
    parser.add_argument(
        "--resume",
        action="store_true",
        help="Resume training from last checkpoint",
    )

    # Stage control
    parser.add_argument(
        "--tile-only",
        action="store_true",
        help="Only run data preparation (extract + tile), then exit",
    )
    parser.add_argument(
        "--train-only",
        action="store_true",
        help="Only run training (requires --data-dir with already-tiled data)",
    )
    parser.add_argument(
        "--export-only",
        type=str,
        metavar="BEST_PT",
        help="Only export the given best.pt to ONNX, then exit",
    )
    parser.add_argument(
        "--stop-pod",
        action="store_true",
        help="Stop the RunPod pod when training completes (saves money). "
        "Uses RUNPOD_POD_ID and RUNPOD_API_KEY from the container environment.",
    )

    args = parser.parse_args()

    # ── Verify network volume ─────────────────────────────────────────────
    if args.output_dir.startswith(VOLUME_DIR) and not Path(VOLUME_DIR).is_dir():
        parser.error(
            f"Network volume not mounted at {VOLUME_DIR}.\n"
            "Deploy the pod with the 'checkbox-training' network volume attached.\n"
            "Training results MUST be saved to persistent storage.\n"
            "To override (not recommended): use --output-dir /workspace/output"
        )

    # ── Validate args ─────────────────────────────────────────────────────
    if not (0.0 <= args.tile_overlap < 1.0):
        parser.error("--tile-overlap must be in [0.0, 1.0)")

    if args.export_only:
        if not Path(args.export_only).exists():
            parser.error(f"--export-only path does not exist: {args.export_only}")
        export(args, args.export_only)
        return

    if args.train_only:
        if not args.data_dir:
            parser.error("--train-only requires --data-dir")
        data_dir = Path(args.data_dir)
        if not data_dir.exists():
            parser.error(f"--data-dir does not exist: {data_dir}")
        stats = train(args, data_dir)

        # Auto-export after training
        if "best_pt" in stats:
            export(args, stats["best_pt"])

        logger.info("\n=== Training Summary ===")
        for k, v in sorted(stats.items()):
            logger.info("  %s: %s", k, v)
        if args.stop_pod:
            _stop_this_pod()
        return

    if not args.tarball and not args.data_dir:
        parser.error("Either --tarball or --data-dir is required")

    # ── Full pipeline ─────────────────────────────────────────────────────
    start = time.time()

    # Stage 1: Prepare data
    if args.data_dir:
        data_dir = Path(args.data_dir)
        if not data_dir.exists():
            parser.error(f"--data-dir does not exist: {data_dir}")
        logger.info("Using pre-extracted data at %s", data_dir)
        _fix_dataset_yaml(data_dir)
    else:
        data_dir = prepare_data(args)

    if args.tile_only:
        elapsed = time.time() - start
        logger.info("Tiling complete in %.1f seconds. Data at: %s", elapsed, data_dir)
        return

    # Stage 2: Train
    stats = train(args, data_dir)

    # Stage 3: Export
    if "best_pt" in stats:
        onnx_path = export(args, stats["best_pt"])
        stats["onnx_path"] = onnx_path

    elapsed = time.time() - start
    logger.info("\n=== Pipeline Complete (%.1f min) ===", elapsed / 60)
    for k, v in sorted(stats.items()):
        logger.info("  %s: %s", k, v)

    if args.stop_pod:
        _stop_this_pod()


def _read_runpod_env() -> Dict[str, str]:
    """Read RunPod env vars, falling back to PID 1's environment.

    SSH sessions don't inherit the container-level RUNPOD_* vars,
    but they're always available in /proc/1/environ.
    """
    result = {}
    for key in ("RUNPOD_POD_ID", "RUNPOD_API_KEY"):
        val = os.environ.get(key)
        if val:
            result[key] = val

    if len(result) == 2:
        return result

    # Fallback: read from PID 1's environment
    try:
        environ = Path("/proc/1/environ").read_bytes()
        for entry in environ.split(b"\0"):
            if b"=" not in entry:
                continue
            k, v = entry.decode("utf-8", errors="replace").split("=", 1)
            if k in ("RUNPOD_POD_ID", "RUNPOD_API_KEY") and k not in result:
                result[k] = v
    except (OSError, PermissionError):
        pass

    return result


def _stop_this_pod():
    """Stop the current RunPod pod via the API.

    Reads RUNPOD_POD_ID and RUNPOD_API_KEY from the environment
    or from /proc/1/environ (container-level vars that SSH doesn't inherit).
    """
    env = _read_runpod_env()
    pod_id = env.get("RUNPOD_POD_ID")
    api_key = env.get("RUNPOD_API_KEY")

    if not pod_id:
        logger.warning("RUNPOD_POD_ID not found — cannot auto-stop pod")
        return

    if not api_key:
        logger.warning("RUNPOD_API_KEY not found — cannot auto-stop pod")
        return

    import urllib.request
    import urllib.error

    query = (
        'mutation { podStop(input: {podId: "%s"}) { id desiredStatus } }' % pod_id
    )
    url = f"https://api.runpod.io/graphql?api_key={api_key}"
    req = urllib.request.Request(
        url,
        data=json.dumps({"query": query}).encode(),
        headers={"Content-Type": "application/json", "User-Agent": "train-cloud/1.0"},
    )

    try:
        with urllib.request.urlopen(req, timeout=30) as resp:
            result = json.loads(resp.read())
            logger.info("Pod %s stop requested: %s", pod_id, result)
    except (urllib.error.URLError, OSError) as e:
        logger.error("Failed to stop pod %s: %s", pod_id, e)


if __name__ == "__main__":
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s %(levelname)s: %(message)s",
        datefmt="%H:%M:%S",
    )
    main()
