#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
run_flywheel.py — one command that closes the human-reinforcement loop.

Stages (select with --stages, comma-separated, or "all"):

  DATA stages (run anywhere, incl. CI — no GPU needed):
    sync     pull predictions + corrections from Supabase → JSONL
             (sync_supabase_feedback.py; needs SUPABASE_URL/-KEY, else skipped)
    ingest   JSONL → YOLO dataset (ingest_feedback.py --download)
    report   dataset/feedback statistics → flywheel_report.json

  TRAIN stages (GPU box / Colab; need `ultralytics` installed):
    train      train_yolov8_seg.py on --data (default: the ingested dataset)
    calibrate  calibrate_temp.py on the new artifacts
    evaluate   evaluate.py on the new artifacts
    export     export_onnx.py → candidate bundle (--candidate-out)

  RELEASE stages:
    gate     quality_gate.py candidate vs deploy/latest (regression = stop)
    promote  promote_bundle.py → deploy/<version> + deploy/latest + registry

Typical uses:
  # CI / laptop: refresh training data + readiness report
  python ml/scripts/run_flywheel.py --stages sync,ingest,report

  # GPU box: full loop
  python ml/scripts/run_flywheel.py --stages all --data ml/datasets/merged/dataset.yaml

Note on training data: the ingested user-feedback set alone is usually too
small/biased for a full retrain. Merge it with the TACO base set first
(ml/scripts/ingest_yolo_dataset.py / augment_external_datasets.py) and pass
the merged dataset.yaml via --data.
"""
from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import subprocess
import sys
from collections import Counter
from pathlib import Path
from typing import Dict, List, Optional

REPO = Path(__file__).resolve().parents[2]
SCRIPTS = REPO / "ml" / "scripts"

DATA_STAGES = ["sync", "ingest", "report"]
TRAIN_STAGES = ["train", "calibrate", "evaluate", "export"]
RELEASE_STAGES = ["gate", "promote"]
ALL_STAGES = DATA_STAGES + TRAIN_STAGES + RELEASE_STAGES


def run(cmd: List[str], allow_fail: bool = False) -> int:
    print(f"\n>>> {' '.join(str(c) for c in cmd)}", flush=True)
    rc = subprocess.call([str(c) for c in cmd])
    if rc != 0 and not allow_fail:
        print(f"FLYWHEEL STOP — stage command failed (rc={rc}).")
        sys.exit(rc)
    return rc


def dataset_stats(ds_dir: Path) -> Dict:
    stats: Dict = {"exists": ds_dir.exists(), "path": str(ds_dir)}
    if not ds_dir.exists():
        return stats
    names_p = ds_dir / "names.json"
    names = []
    if names_p.exists():
        try:
            names = json.loads(names_p.read_text(encoding="utf-8"))
        except Exception:
            pass
    class_counts: Counter = Counter()
    for split in ("train", "val"):
        imgs = list((ds_dir / "images" / split).glob("*")) if (ds_dir / "images" / split).exists() else []
        labels = list((ds_dir / "labels" / split).glob("*.txt")) if (ds_dir / "labels" / split).exists() else []
        stats[split] = {"images": len(imgs), "labels": len(labels)}
        for lp in labels:
            try:
                for line in lp.read_text(encoding="utf-8").splitlines():
                    parts = line.split()
                    if parts:
                        cid = int(parts[0])
                        cname = names[cid] if 0 <= cid < len(names) else str(cid)
                        class_counts[cname] += 1
            except Exception:
                continue
    stats["class_counts"] = dict(class_counts)
    stats["classes"] = names
    return stats


def newest_artifacts_dir(root: Path) -> Optional[Path]:
    if not root.exists():
        return None
    candidates = [d for d in root.iterdir() if d.is_dir()]
    return max(candidates, key=lambda d: d.stat().st_mtime) if candidates else None


def main(argv=None) -> int:
    ap = argparse.ArgumentParser(description="Run the Alami Vision human-reinforcement flywheel.")
    ap.add_argument("--stages", default="sync,ingest,report",
                    help=f"comma-separated from {ALL_STAGES} or 'all'")
    ap.add_argument("--sync-out", default="feedback_logs_sync", type=Path)
    ap.add_argument("--dataset-out", default="ml/datasets/alami_user", type=Path)
    ap.add_argument("--data", default=None, type=Path,
                    help="dataset.yaml for training (default: <dataset-out>/dataset.yaml)")
    ap.add_argument("--artifacts-root", default="artifacts", type=Path)
    ap.add_argument("--candidate-out", default="deploy_candidate", type=Path)
    ap.add_argument("--max-samples", type=int, default=5000)
    ap.add_argument("--report-out", default="flywheel_report.json", type=Path)
    args = ap.parse_args(argv)

    stages = ALL_STAGES if args.stages.strip() == "all" else [s.strip() for s in args.stages.split(",") if s.strip()]
    unknown = [s for s in stages if s not in ALL_STAGES]
    if unknown:
        print(f"Unknown stages: {unknown}. Valid: {ALL_STAGES}")
        return 2

    py = sys.executable
    report: Dict = {"stages_run": stages}

    # ---------------- DATA ----------------
    if "sync" in stages:
        if os.environ.get("SUPABASE_URL") and os.environ.get("SUPABASE_SERVICE_ROLE_KEY"):
            run([py, SCRIPTS / "sync_supabase_feedback.py", "--out", args.sync_out])
            stats_p = args.sync_out / "sync_stats.json"
            if stats_p.exists():
                report["sync"] = json.loads(stats_p.read_text(encoding="utf-8"))
        else:
            print("sync: SUPABASE_URL/SUPABASE_SERVICE_ROLE_KEY not set — skipping "
                  "(will use local feedback_logs/ if present).")
            report["sync"] = {"skipped": "no supabase credentials"}

    if "ingest" in stages:
        pred_file = args.sync_out / "predictions.jsonl"
        fb_file = args.sync_out / "feedback.jsonl"
        if not pred_file.exists():
            pred_file = Path("feedback_logs/predictions.jsonl")
            fb_file = Path("feedback_logs/feedback.jsonl")
        if pred_file.exists():
            run([py, SCRIPTS / "ingest_feedback.py",
                 "--predictions", pred_file, "--feedback", fb_file,
                 "--out", args.dataset_out, "--max", str(args.max_samples),
                 "--names", "deploy/latest/names.json", "--download"])
        else:
            print("ingest: no predictions.jsonl found (neither synced nor local) — skipping.")
            report["ingest"] = {"skipped": "no source jsonl"}

    if "report" in stages:
        report["dataset"] = dataset_stats(args.dataset_out)
        args.report_out.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
        print(f"\nFlywheel report -> {args.report_out}")
        print(json.dumps(report.get("dataset", {}), indent=2)[:2000])

    # ---------------- TRAIN ----------------
    needs_training = any(s in stages for s in TRAIN_STAGES)
    artifacts_dir: Optional[Path] = None
    if needs_training:
        try:
            import ultralytics  # noqa: F401
        except ImportError:
            print("TRAIN stages need `ultralytics` (pip install ultralytics). "
                  "Run these stages on a GPU box/Colab — data stages already completed.")
            return 3

    data_yaml = args.data or (args.dataset_out / "dataset.yaml")

    if "train" in stages:
        run([py, SCRIPTS / "train_yolov8_seg.py", "--data", data_yaml,
             "--artifacts_root", args.artifacts_root])
    if any(s in stages for s in ("calibrate", "evaluate", "export")):
        artifacts_dir = newest_artifacts_dir(args.artifacts_root)
        if artifacts_dir is None:
            print(f"No artifacts found under {args.artifacts_root} — did training run?")
            return 3
        print(f"Using artifacts: {artifacts_dir}")
    if "calibrate" in stages:
        run([py, SCRIPTS / "calibrate_temp.py", "--artifacts_dir", artifacts_dir,
             "--dataset_yaml", data_yaml])
    if "evaluate" in stages:
        run([py, SCRIPTS / "evaluate.py", "--artifacts_dir", artifacts_dir,
             "--dataset_yaml", data_yaml])
    if "export" in stages:
        # A stale gate report must never survive a re-export (the hash binding
        # in promote_bundle would catch it, but don't even leave it around).
        (args.candidate_out / "gate_report.json").unlink(missing_ok=True)
        run([py, SCRIPTS / "export_onnx.py", "--artifacts_dir", artifacts_dir,
             "--out_dir", args.candidate_out, "--simplify"])

    # ---------------- RELEASE ----------------
    baseline = REPO / "deploy" / "latest"  # absolute: gate must not depend on cwd
    if "gate" in stages:
        rc = run([py, SCRIPTS / "quality_gate.py", "--candidate", args.candidate_out,
                  "--baseline", baseline], allow_fail=True)
        if rc != 0:
            print("Gate failed — candidate NOT promoted. Inspect gate_report.json.")
            return 1
    if "promote" in stages:
        # Traceable version name: artifacts dir if we trained in this run,
        # else the candidate's model_card lineage, else a timestamp.
        version = artifacts_dir.name if artifacts_dir is not None else None
        if not version:
            try:
                card = json.loads((args.candidate_out / "model_card.json").read_text(encoding="utf-8"))
                version = Path(card.get("artifacts_dir", "")).name or None
            except Exception:
                version = None
        if not version:
            version = dt.datetime.now(dt.timezone.utc).strftime("v%Y%m%d-%H%M%S")
        run([py, SCRIPTS / "promote_bundle.py", "--candidate", args.candidate_out,
             "--version", version,
             "--deploy-root", REPO / "deploy",
             "--registry", REPO / "model_registry" / "registry.json"])

    print("\nFlywheel done.")
    return 0


if __name__ == "__main__":
    sys.exit(main())