#!/usr/bin/env python3
"""Stage 04 — Extract frames from GoPro videos.

For each MP4 in /mnt/ssd/<mission>/raw_data/medias/videos/GP*-AUV*/
- Skip files < 2MB (placeholders)
- Auto-trim hors-eau: sample frames at start/end, detect non-blue/green pixels
- ffmpeg fps=1, scale=518:294, q:v=3
- Output: ~/cosma-pipeline/data/<mission>/frames/<AUV>/<segment>/frame_XXXXX.jpg
- Skip if output dir exists and has >= expected frames
- Log to SQLite state.db

Usage:
    python3 04_frame_extract.py --mission 20260505-Lepradet
    python3 04_frame_extract.py --video /mnt/ssd/.../GP1-AUV210/GX019837.MP4
"""
from __future__ import annotations

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

import cv2
import numpy as np

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

QC_SAMPLE_RATE = int(os.environ.get("COSMA_QC_SAMPLE_RATE", "5"))
QC_BOTTOM_OK_PCT = float(os.environ.get("COSMA_QC_BOTTOM_OK_PCT", "50"))

PIPELINE_BASE = Path(os.environ.get("COSMA_PIPELINE_BASE", "/home/cosma/cosma-pipeline"))
SSD_BASE = Path(os.environ.get("COSMA_SSD_BASE", "/mnt/ssd"))
MIN_VIDEO_SIZE_MB = 2.0


def is_underwater_frame(frame_bgr: np.ndarray, threshold: float = 0.6) -> bool:
    """Return True if frame looks like underwater footage (dominant blue/green).
    Hors-eau: R > G-5 AND R > B-5 (dry/air dominant).
    Underwater: blue or green channel dominant.
    """
    b, g, r = cv2.split(frame_bgr.astype(np.float32))
    mean_r = float(np.mean(r))
    mean_g = float(np.mean(g))
    mean_b = float(np.mean(b))
    # Not underwater: red dominates
    if mean_r > mean_g - 5 and mean_r > mean_b - 5:
        return False
    return True


def detect_water_range(video_path: Path, sample_count: int = 10) -> tuple[float, float]:
    """Detect start/end times of underwater portion by sampling frames.
    Returns (start_s, end_s). Returns (0, duration) if uncertain.
    """
    cap = cv2.VideoCapture(str(video_path))
    if not cap.isOpened():
        return 0.0, -1.0

    fps = cap.get(cv2.CAP_PROP_FPS) or 25.0
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration_s = total_frames / fps if fps > 0 else 0

    # Sample frames: first 20% and last 20%
    probe_times_start = [duration_s * i / (sample_count * 5) for i in range(sample_count)]
    probe_times_end = [duration_s * (1 - i / (sample_count * 5)) for i in range(sample_count)]

    # Find first underwater frame from start
    start_s = 0.0
    for t in probe_times_start:
        cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000)
        ret, frame = cap.read()
        if ret and is_underwater_frame(frame):
            start_s = max(0.0, t - 2.0)
            break

    # Find last underwater frame from end
    end_s = duration_s
    for t in sorted(probe_times_end, reverse=True):
        cap.set(cv2.CAP_PROP_POS_MSEC, t * 1000)
        ret, frame = cap.read()
        if ret and is_underwater_frame(frame):
            end_s = min(duration_s, t + 2.0)
            break

    cap.release()
    return start_s, end_s


def get_video_duration(video_path: Path) -> float:
    """Get video duration in seconds via ffprobe."""
    cmd = [
        "ffprobe", "-v", "quiet", "-print_format", "json",
        "-show_streams", str(video_path)
    ]
    try:
        r = subprocess.run(cmd, capture_output=True, text=True, timeout=30)
        data = json.loads(r.stdout)
        for stream in data.get("streams", []):
            dur = float(stream.get("duration", 0))
            if dur > 0:
                return dur
    except Exception:
        pass
    return 0.0


def extract_frames(video_path: Path, out_dir: Path, fps: int = 1,
                   scale: str = "518:294", quality: int = 3,
                   start_s: float = 0.0, end_s: float = -1.0) -> dict:
    """Run ffmpeg to extract frames. Returns metrics dict."""
    out_dir.mkdir(parents=True, exist_ok=True)

    # Build ffmpeg args
    cmd = ["ffmpeg", "-y", "-loglevel", "error"]
    cmd += ["-ss", str(start_s), "-i", str(video_path)]
    if end_s > 0 and end_s > start_s:
        cmd += ["-t", str(end_s - start_s)]
    cmd += [
        "-vf", f"fps={fps},scale={scale}",
        "-q:v", str(quality),
        str(out_dir / "frame_%05d.jpg"),
    ]

    t0 = time.time()
    result = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
    elapsed = time.time() - t0

    frames = sorted(out_dir.glob("frame_*.jpg"))
    n_frames = len(frames)

    metrics = {
        "video": str(video_path),
        "out_dir": str(out_dir),
        "n_frames": n_frames,
        "elapsed_s": round(elapsed, 1),
        "returncode": result.returncode,
        "start_s": start_s,
        "end_s": end_s,
    }
    if result.returncode != 0:
        metrics["error"] = result.stderr[-500:]
        print(f"  [04] ffmpeg error for {video_path.name}: {result.stderr[-200:]}")
    else:
        print(f"  [04] {video_path.name}: {n_frames} frames in {elapsed:.1f}s")
        # Score a subsample for QC
        qc = qc_segment(out_dir, sample_rate=QC_SAMPLE_RATE)
        if qc:
            metrics.update(qc)

    return metrics


def qc_segment(frames_dir: Path, sample_rate: int = 5) -> dict | None:
    """Sample 1/sample_rate frames, score each, write qc.json, return aggregate."""
    frames = sorted(frames_dir.glob("frame_*.jpg"))
    if not frames:
        return None
    sampled = frames[::max(1, sample_rate)]
    per_frame = []
    for f in sampled:
        s = score_image_file(f)
        if s is not None:
            per_frame.append(s)
    if not per_frame:
        return None
    agg = qc_aggregate(per_frame)
    qc_payload = {
        "frames_in_dir": len(frames),
        "frames_sampled": len(per_frame),
        "sample_rate": sample_rate,
        **agg,
        "per_frame": per_frame,
    }
    try:
        (frames_dir / "qc.json").write_text(json.dumps(qc_payload, indent=2))
    except Exception as e:
        print(f"  [04] qc.json write failed: {e}")
    print(
        f"  [04]   QC: bottom_visible={agg['bottom_visible_pct']}% "
        f"(b={agg['frames_bottom_visible']} ooo={agg['frames_out_of_water']} "
        f"turb={agg['frames_turbid']} nob={agg['frames_water_no_bottom']})"
    )
    return agg


def process_video(video_path: Path, auv_id: str, mission_name: str) -> dict:
    """Process one video file. Returns metrics."""
    size_mb = video_path.stat().st_size / (1024 * 1024)
    if size_mb < MIN_VIDEO_SIZE_MB:
        print(f"  [04] Skip {video_path.name} ({size_mb:.1f}MB < {MIN_VIDEO_SIZE_MB}MB)")
        return {"video": str(video_path), "skipped": True, "reason": "placeholder"}

    segment = video_path.stem
    out_dir = PIPELINE_BASE / "data" / mission_name / "frames" / auv_id / segment

    # Check if already done
    if out_dir.exists():
        existing = list(out_dir.glob("frame_*.jpg"))
        duration_s = get_video_duration(video_path)
        expected = max(1, int(duration_s) - 10) if duration_s > 0 else 1
        if len(existing) >= expected:
            print(f"  [04] {video_path.name}: already done ({len(existing)} frames), skip")
            cached_m: dict = {"video": str(video_path), "n_frames": len(existing), "cached": True,
                              "out_dir": str(out_dir)}
            # Re-run QC if qc.json is missing (idempotent enrichment)
            if not (out_dir / "qc.json").exists():
                qc = qc_segment(out_dir, sample_rate=QC_SAMPLE_RATE)
                if qc:
                    cached_m.update(qc)
            else:
                try:
                    cached_qc = json.loads((out_dir / "qc.json").read_text())
                    for k in (
                        "frames_total", "frames_bottom_visible", "frames_out_of_water",
                        "frames_turbid", "frames_water_no_bottom", "bottom_visible_pct",
                    ):
                        if k in cached_qc:
                            cached_m[k] = cached_qc[k]
                except Exception:
                    pass
            return cached_m

    print(f"  [04] {video_path.name} ({size_mb:.0f}MB): detecting water range...")
    start_s, end_s = detect_water_range(video_path)
    print(f"  [04]   water range: {start_s:.1f}s → {end_s:.1f}s")

    return extract_frames(video_path, out_dir, start_s=start_s, end_s=end_s)


def find_auv_videos(mission_path: Path) -> dict[str, list[Path]]:
    """Find all MP4 files per AUV in medias/videos/GP*-AUV*/."""
    videos_root = mission_path / "raw_data/medias/videos"
    result: dict[str, list[Path]] = {}

    for gopro_dir in sorted(videos_root.glob("GP*-AUV*")):
        # Extract AUV ID from dir name: GP1-AUV210 -> AUV210
        parts = gopro_dir.name.split("-")
        if len(parts) >= 2:
            auv_id = parts[1]
            mp4_files = [f for f in sorted(gopro_dir.glob("GX*.MP4"))
                         if f.stat().st_size / (1024 * 1024) >= MIN_VIDEO_SIZE_MB]
            if mp4_files:
                if auv_id not in result:
                    result[auv_id] = []
                result[auv_id].extend(mp4_files)

    return result


def process_mission(mission_name: str) -> list[dict]:
    mission_path = SSD_BASE / mission_name
    auv_videos = find_auv_videos(mission_path)
    print(f"[04] Found AUVs: {list(auv_videos.keys())}")

    all_metrics = []
    init_db()

    for auv_id, videos in sorted(auv_videos.items()):
        print(f"[04] === {auv_id}: {len(videos)} videos ===")
        for video_path in videos:
            t0 = time.time()
            m = process_video(video_path, auv_id, mission_name)
            m["auv_id"] = auv_id
            all_metrics.append(m)

            if not m.get("skipped"):
                with get_conn() as conn:
                    mission_row = conn.execute(
                        "SELECT id FROM missions WHERE name=?", (mission_name,)
                    ).fetchone()
                    if mission_row:
                        bottom_pct = m.get("bottom_visible_pct")
                        if m.get("returncode", 0) != 0:
                            job_status = "error"
                        elif bottom_pct is not None and bottom_pct < QC_BOTTOM_OK_PCT:
                            job_status = "degraded"
                        else:
                            job_status = "done"
                        job_id = upsert_job(
                            conn, mission_row["id"], auv_id,
                            video_path.stem, "04_frame_extract",
                            status=job_status,
                            output_path=m.get("out_dir", ""),
                            error_msg=(
                                f"bottom_visible_pct={bottom_pct}% <{QC_BOTTOM_OK_PCT}%"
                                if job_status == "degraded" else None
                            ),
                        )
                        if not m.get("cached"):
                            record_metric(conn, job_id, "frames_extracted",
                                          value=m.get("n_frames", 0),
                                          pass_fail="pass" if m.get("n_frames", 0) > 0 else "fail")
                            record_metric(conn, job_id, "extract_time_s",
                                          value=m.get("elapsed_s", 0))
                        # Always record QC metrics (so cached frames also get scored history)
                        for k in (
                            "frames_total", "frames_bottom_visible", "frames_out_of_water",
                            "frames_turbid", "frames_water_no_bottom",
                        ):
                            if k in m:
                                record_metric(conn, job_id, k, value=float(m[k]))
                        if bottom_pct is not None:
                            record_metric(
                                conn, job_id, "bottom_visible_pct",
                                value=float(bottom_pct),
                                pass_fail="pass" if bottom_pct >= QC_BOTTOM_OK_PCT else "degraded",
                            )

    return all_metrics


def main():
    ap = argparse.ArgumentParser(description="Stage 04 — Extract frames from GoPro videos")
    ap.add_argument("--mission", type=str, help="Mission name (e.g. 20260505-Lepradet)")
    ap.add_argument("--video", type=Path, help="Single video path")
    ap.add_argument("--auv", type=str, default="UNKNOWN", help="AUV ID for single video mode")
    args = ap.parse_args()

    if args.video:
        mission_name = args.mission or "unknown"
        m = process_video(args.video, args.auv, mission_name)
        print(f"\nResult: {m}")
    elif args.mission:
        metrics = process_mission(args.mission)
        print("\n=== Stage 04 summary ===")
        total_frames = sum(m.get("n_frames", 0) for m in metrics if not m.get("skipped"))
        skipped = sum(1 for m in metrics if m.get("skipped"))
        print(f"Total frames: {total_frames}, skipped: {skipped}")
    else:
        ap.print_help()
        sys.exit(1)


if __name__ == "__main__":
    main()
