#!/usr/bin/env python3
"""
Sync trained .tflite models from Colab output into the sample app's assets dir.

The Colab notebook writes trained models to ``collab/output/<run-id>/*.tflite``
(via Drive sync). This script copies those models into the sample app's
Android assets directory so ``DiscoverModels.android.kt`` will pick them up
on the next build.

Usage::

    tools/sync_models.py                    # sync the latest run (default)
    tools/sync_models.py --list             # list available runs newest-first
    tools/sync_models.py --run <run-id>     # sync a specific run
    tools/sync_models.py --clean            # remove previously-synced models

Synced files are renamed ``<run-id>__<model-base>.tflite`` (double underscore).
The double underscore is deliberate: ``DiscoverModels.android.kt`` replaces
single underscores with spaces in the picker, so ``__`` becomes a clear visual
break between the run-id and the model name in the dropdown.

``--clean`` preserves the three baseline models bundled with the repo.
"""

from __future__ import annotations

import argparse
import re
import shutil
import sys
from pathlib import Path

REPO_ROOT = Path(__file__).resolve().parent.parent
COLAB_OUTPUT_DIR = REPO_ROOT / "collab" / "output"
ASSETS_DIR = REPO_ROOT / "sample" / "composeApp" / "src" / "androidMain" / "assets"

# Baseline models that ship with the repo. --clean preserves these; everything
# else under ASSETS_DIR with a .tflite extension is treated as a synced file
# and removed by --clean.
BASELINE_MODELS = frozenset(
    {
        "yolo11n_dataset_dataset.tflite",
        "yolo11n_su_416.tflite",
        "yolov10n_float16.tflite",
    }
)

SAFE_NAME = re.compile(r"[^A-Za-z0-9._-]")


def sanitize(name: str) -> str:
    """Replace anything that isn't safe for an Android asset filename."""
    return SAFE_NAME.sub("_", name).strip("_") or "unnamed"


def list_runs() -> list[Path]:
    """Return run-id directories under collab/output, newest first by mtime."""
    if not COLAB_OUTPUT_DIR.exists():
        return []
    runs = [p for p in COLAB_OUTPUT_DIR.iterdir() if p.is_dir()]
    runs.sort(key=lambda p: p.stat().st_mtime, reverse=True)
    return runs


def cmd_list() -> int:
    runs = list_runs()
    if not runs:
        print(f"(no runs found under {COLAB_OUTPUT_DIR.relative_to(REPO_ROOT)})")
        return 0
    print(f"Runs under {COLAB_OUTPUT_DIR.relative_to(REPO_ROOT)} (newest first):")
    for run in runs:
        tflites = sorted(run.glob("*.tflite"))
        if tflites:
            tflite_summary = ", ".join(p.name for p in tflites)
        else:
            tflite_summary = "(no .tflite files)"
        print(f"  {run.name}  —  {tflite_summary}")
    return 0


def cmd_sync(run_id: str) -> int:
    run_dir = COLAB_OUTPUT_DIR / run_id
    if not run_dir.is_dir():
        print(f"error: no such run directory: {run_dir}", file=sys.stderr)
        available = [r.name for r in list_runs()]
        if available:
            print(f"  available: {', '.join(available)}", file=sys.stderr)
        return 2

    tflites = sorted(run_dir.glob("*.tflite"))
    if not tflites:
        print(f"error: no .tflite files in {run_dir}", file=sys.stderr)
        return 2

    if not ASSETS_DIR.exists():
        print(f"error: assets dir does not exist: {ASSETS_DIR}", file=sys.stderr)
        return 2

    safe_run = sanitize(run_id)
    print(f"Syncing run '{run_id}' ({len(tflites)} model(s)) -> {ASSETS_DIR.relative_to(REPO_ROOT)}")
    for src in tflites:
        base = sanitize(src.stem)
        dest_name = f"{safe_run}__{base}.tflite"
        if dest_name in BASELINE_MODELS:
            print(f"  skip {src.name} (would shadow baseline {dest_name})")
            continue
        dest = ASSETS_DIR / dest_name
        shutil.copy2(src, dest)
        size_mb = dest.stat().st_size / (1024 * 1024)
        print(f"  copied {src.name} -> {dest.name} ({size_mb:.1f} MB)")
    return 0


def cmd_clean() -> int:
    if not ASSETS_DIR.exists():
        print(f"error: assets dir does not exist: {ASSETS_DIR}", file=sys.stderr)
        return 2
    removed = 0
    for f in sorted(ASSETS_DIR.glob("*.tflite")):
        if f.name in BASELINE_MODELS:
            continue
        f.unlink()
        print(f"  removed {f.name}")
        removed += 1
    if removed == 0:
        print("(nothing to clean — only baseline models present)")
    else:
        print(f"removed {removed} synced model(s); baselines preserved")
    return 0


def cmd_latest() -> int:
    runs = list_runs()
    if not runs:
        print(f"error: no runs under {COLAB_OUTPUT_DIR}", file=sys.stderr)
        return 2
    return cmd_sync(runs[0].name)


def main(argv: list[str]) -> int:
    parser = argparse.ArgumentParser(
        description="Sync Colab-trained .tflite models into the sample app's assets dir.",
    )
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--list", action="store_true", help="list available runs newest-first")
    group.add_argument("--run", metavar="RUN_ID", help="sync the named run from collab/output/")
    group.add_argument("--clean", action="store_true", help="remove previously-synced models (preserve baselines)")
    args = parser.parse_args(argv)

    if args.list:
        return cmd_list()
    if args.clean:
        return cmd_clean()
    if args.run:
        return cmd_sync(args.run)
    # Default: sync the latest run.
    return cmd_latest()


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))
