#!/usr/bin/env python3
"""Stage 05 — Run lingbot-map inference on extracted frames.

Args:
    --frames-dir <path>   Directory with frame_*.jpg (or parent with AUV subdirs)
    --worker <auto|.84|.87>   GPU worker selection
    --mission <name>      Mission name for output paths

Workers:
  .84: /root/ai-video/lingbot-map/.venv/bin/python demo.py ...
  .87: /home/floppyrj45/ai-video/lingbot-map/.venv/bin/python demo.py ...

Auto: pick by lowest GPU memory usage (nvidia-smi via SSH).

Flow:
1. rsync frames .83 → worker /root/cosma-frames-tmp/ (or /home/floppyrj45/)
2. SSH launch demo.py with windowed mode (window=64, overlap=16)
3. Retrieve PLY + NPZ → .83 ~/cosma-pipeline/data/<mission>/ply/<AUV>/<segment>.{ply,npz}
4. Cleanup worker temp dir
5. Log to SQLite: duration, GPU peak mem, nb points in PLY

Usage:
    python3 05_inference.py --frames-dir ~/cosma-pipeline/data/20260505-Lepradet/frames/AUV210/GX019837 --worker auto --mission 20260505-Lepradet
"""
from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent.parent))
from orchestrator.db import init_db, get_conn, upsert_job, record_metric, now_iso

PIPELINE_BASE = Path(os.environ.get("COSMA_PIPELINE_BASE", "/home/cosma/cosma-pipeline"))

WORKERS = {
    ".84": {
        "host": "192.168.0.84",
        "user": "root",
        "ai_dir": "/root/ai-video/lingbot-map",
        "venv": "/root/ai-video/lingbot-map/.venv/bin/python",
        "tmp_dir": "/root/cosma-frames-tmp",
    },
    ".87": {
        "host": "192.168.0.87",
        "user": "floppyrj45",
        "ai_dir": "/home/floppyrj45/ai-video/lingbot-map",
        "venv": "/home/floppyrj45/ai-video/lingbot-map/.venv/bin/python",
        "tmp_dir": "/home/floppyrj45/cosma-frames-tmp",
    },
}


def get_gpu_mem_used(worker_key: str) -> int:
    """Return GPU memory used in MB via SSH nvidia-smi. Returns 99999 on error."""
    w = WORKERS[worker_key]
    cmd = [
        "ssh", "-o", "StrictHostKeyChecking=no", "-o", "ConnectTimeout=5",
        f"{w['user']}@{w['host']}",
        "nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits 2>/dev/null | head -1"
    ]
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=10)
        return int(r.stdout.strip())
    except Exception:
        return 99999


def pick_worker() -> str:
    """Auto-select worker with lowest GPU memory usage."""
    best = None
    best_mem = 99999
    for key in WORKERS:
        mem = get_gpu_mem_used(key)
        print(f"  [05] Worker {key}: GPU mem={mem}MB")
        if mem < best_mem:
            best_mem = mem
            best = key
    if best is None:
        raise RuntimeError("No GPU worker available")
    print(f"  [05] Selected worker {best}")
    return best


def count_ply_points(ply_path: Path) -> int:
    """Count vertex count in PLY file header."""
    try:
        with open(ply_path, "rb") as f:
            for _ in range(30):
                line = f.readline().decode("ascii", errors="ignore").strip()
                if line.startswith("element vertex"):
                    return int(line.split()[-1])
    except Exception:
        pass
    return 0


def run_inference(frames_dir: Path, worker_key: str, mission_name: str,
                  auv_id: str, segment: str) -> dict:
    """Run lingbot-map on one segment. Returns metrics."""
    w = WORKERS[worker_key]
    host = w["host"]
    user = w["user"]
    ssh_target = f"{user}@{host}"
    worker_frames = f"{w['tmp_dir']}/{mission_name}/{auv_id}/{segment}"
    ply_remote = f"{w['tmp_dir']}/{mission_name}/{auv_id}/{segment}.ply"
    npz_remote = f"{w['tmp_dir']}/{mission_name}/{auv_id}/{segment}.npz"

    out_dir = PIPELINE_BASE / "data" / mission_name / "ply" / auv_id
    out_dir.mkdir(parents=True, exist_ok=True)
    out_ply = out_dir / f"{segment}.ply"
    out_npz = out_dir / f"{segment}.npz"

    if out_ply.exists() and out_ply.stat().st_size > 1000:
        n_pts = count_ply_points(out_ply)
        print(f"  [05] {auv_id}/{segment}: cached PLY ({n_pts} pts)")
        return {"cached": True, "ply": str(out_ply), "n_points": n_pts}

    metrics = {
        "auv_id": auv_id,
        "segment": segment,
        "worker": worker_key,
        "status": "ok",
    }

    # Step 1: create remote temp dir + rsync frames
    print(f"  [05] rsync {frames_dir} → {ssh_target}:{worker_frames}...")
    subprocess.run(
        ["ssh", "-o", "StrictHostKeyChecking=no", ssh_target,
         f"mkdir -p {worker_frames}"],
        check=True, timeout=15,
    )
    r = subprocess.run(
        ["rsync", "-az", "--delete",
         str(frames_dir) + "/",
         f"{ssh_target}:{worker_frames}/"],
        capture_output=True, text=True, timeout=600,
    )
    if r.returncode != 0:
        metrics["status"] = "error"
        metrics["error"] = f"rsync failed: {r.stderr[-200:]}"
        return metrics
    print(f"  [05] rsync done")

    # Step 2: build demo.py command
    checkpoint = f"{w['ai_dir']}/checkpoints/lingbot-map/lingbot-map.pt"
    demo_cmd = (
        f"cd {w['ai_dir']} && "
        f"{w['venv']} demo.py "
        f"--model_path {checkpoint} "
        f"--image_folder {worker_frames} "
        f"--mode windowed "
        f"--window_size 64 "
        f"--overlap_size 16 "
        f"--save_ply {ply_remote} "
        f"--save_poses {npz_remote} "
        f"--use_sdpa "
        f"2>&1"
    )

    print(f"  [05] Launching inference on {host}...")
    t0 = time.time()
    r = subprocess.run(
        ["ssh", "-o", "StrictHostKeyChecking=no", ssh_target, demo_cmd],
        capture_output=True, text=True, timeout=7200,  # 2h max
    )
    elapsed = time.time() - t0
    metrics["inference_s"] = round(elapsed, 1)

    if r.returncode != 0:
        metrics["status"] = "error"
        metrics["error"] = r.stdout[-500:] + r.stderr[-200:]
        print(f"  [05] inference error: {metrics['error'][-200:]}")
        return metrics

    print(f"  [05] Inference done in {elapsed:.1f}s")

    # Step 3: GPU peak mem from nvidia-smi log (best-effort parse)
    gpu_mem_line = [l for l in r.stdout.split("\n") if "MiB" in l]
    metrics["gpu_peak_mb"] = get_gpu_mem_used(worker_key)

    # Step 4: rsync PLY + NPZ back
    print(f"  [05] Retrieving PLY + NPZ...")
    for remote, local in [(ply_remote, out_ply), (npz_remote, out_npz)]:
        r2 = subprocess.run(
            ["rsync", "-az", f"{ssh_target}:{remote}", str(local)],
            capture_output=True, text=True, timeout=120,
        )
        if r2.returncode != 0:
            print(f"  [05] Warning: rsync back failed for {remote}: {r2.stderr[-100:]}")

    # Step 5: cleanup worker
    subprocess.run(
        ["ssh", "-o", "StrictHostKeyChecking=no", ssh_target,
         f"rm -rf {worker_frames} {ply_remote} {npz_remote}"],
        timeout=30,
    )

    # Count PLY points
    n_pts = count_ply_points(out_ply) if out_ply.exists() else 0
    metrics["n_points"] = n_pts
    metrics["ply"] = str(out_ply)
    print(f"  [05] PLY: {n_pts} points → {out_ply}")

    return metrics


def process_frames_dir(frames_dir: Path, worker_key: str, mission_name: str) -> list[dict]:
    """Process a directory of frames (single segment or AUV tree)."""
    # Detect if frames_dir contains frame_*.jpg directly or subdirs
    direct_frames = list(frames_dir.glob("frame_*.jpg"))

    if direct_frames:
        # Single segment
        parts = frames_dir.parts
        auv_id = frames_dir.parent.name if len(parts) >= 2 else "UNKNOWN"
        segment = frames_dir.name
        return [run_inference(frames_dir, worker_key, mission_name, auv_id, segment)]

    # Tree: frames_dir/<AUV>/<segment>/frame_*.jpg
    all_metrics = []
    for auv_dir in sorted(frames_dir.iterdir()):
        if not auv_dir.is_dir():
            continue
        auv_id = auv_dir.name
        for seg_dir in sorted(auv_dir.iterdir()):
            if not seg_dir.is_dir():
                continue
            frames = list(seg_dir.glob("frame_*.jpg"))
            if not frames:
                continue
            print(f"\n[05] === {auv_id}/{seg_dir.name}: {len(frames)} frames ===")
            m = run_inference(seg_dir, worker_key, mission_name, auv_id, seg_dir.name)
            all_metrics.append(m)

            init_db()
            with get_conn() as conn:
                mission_row = conn.execute(
                    "SELECT id FROM missions WHERE name=?", (mission_name,)
                ).fetchone()
                if mission_row and not m.get("cached"):
                    job_id = upsert_job(
                        conn, mission_row["id"], auv_id, seg_dir.name, "05_inference",
                        status="done" if m.get("status") == "ok" else m.get("status", "error"),
                        output_path=m.get("ply", ""),
                    )
                    record_metric(conn, job_id, "ply_points", value=m.get("n_points", 0),
                                  pass_fail="pass" if m.get("n_points", 0) > 100 else "fail")
                    if "inference_s" in m:
                        record_metric(conn, job_id, "inference_s", value=m["inference_s"])
                    if "gpu_peak_mb" in m:
                        record_metric(conn, job_id, "gpu_peak_mb", value=m["gpu_peak_mb"])

    return all_metrics


def main():
    ap = argparse.ArgumentParser(description="Stage 05 — lingbot-map inference")
    ap.add_argument("--frames-dir", type=Path, required=True,
                    help="Frames dir (single segment or AUV tree)")
    ap.add_argument("--worker", type=str, default="auto",
                    choices=["auto", ".84", ".87"])
    ap.add_argument("--mission", type=str, required=True,
                    help="Mission name (e.g. 20260505-Lepradet)")
    args = ap.parse_args()

    worker = args.worker
    if worker == "auto":
        worker = pick_worker()

    metrics = process_frames_dir(args.frames_dir, worker, args.mission)

    print("\n=== Stage 05 summary ===")
    total_pts = sum(m.get("n_points", 0) for m in metrics)
    ok = sum(1 for m in metrics if m.get("status") == "ok" or m.get("cached"))
    print(f"  Segments OK: {ok}/{len(metrics)}, total PLY points: {total_pts}")
    for m in metrics:
        print(f"  {m.get('auv_id','?')}/{m.get('segment','?')}: "
              f"{m.get('n_points',0)} pts "
              f"[{m.get('status','cached' if m.get('cached') else '?')}]")


if __name__ == "__main__":
    main()
