#!/usr/bin/env python3
"""Stage 04b — Trim out-of-water (hors-eau) head/tail frames from already-extracted segments.

Ports the sustained-run trim logic from cosma-qc/scripts/dispatcher.py (_AUTO_TRIM_SCRIPT,
trim_above_water_prefix) into the new cosma-pipeline pipeline. Re-runs frame QC scoring
on the trimmed set and updates state.db (jobs.status + metrics).

Usage:
    python3 04b_trim_water.py --mission 20260505-Lepradet
    python3 04b_trim_water.py --mission 20260505-Lepradet --auv AUV210 --segment GX019837
    python3 04b_trim_water.py --mission 20260505-Lepradet --dry-run

Safety:
    - Skips segments where ffmpeg is still running on the frames dir (extraction in progress).
    - Skips segments with a queued/running 05_inference job in state.db.
    - Skips segments whose frame count is not stable over a 5s window.
    - Never deletes all frames (sanity floor: keep everything if trim would empty the dir).
"""
from __future__ import annotations

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

import cv2

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

PIPELINE_BASE = Path(os.environ.get("COSMA_PIPELINE_BASE", "/home/cosma/cosma-pipeline"))
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"))
NEED_STREAK = 10  # consecutive underwater frames required to lock start/end


# -----------------------------------------------------------------------------
# Trim logic (ported verbatim from dispatcher._AUTO_TRIM_SCRIPT)
# -----------------------------------------------------------------------------

def is_underwater(path: Path) -> bool | None:
    img = cv2.imread(str(path), cv2.IMREAD_REDUCED_COLOR_4)
    if img is None:
        return None
    b, g, r = [float(c) for c in cv2.mean(img)[:3]]
    # Red is absorbed by water → R < G AND R < B on underwater shots.
    return r < g - 5 and r < b - 5


def trim_segment(frames_dir: Path, dry_run: bool = False) -> tuple[int, int, int]:
    """Delete leading and trailing out-of-water frames.
    Returns (head_removed, tail_removed, remaining).
    """
    paths = sorted(frames_dir.glob("frame_*.jpg"))
    if not paths:
        return (0, 0, 0)

    # Scan from start
    start = 0
    streak = 0
    for i, p in enumerate(paths):
        uw = is_underwater(p)
        if uw is None:
            continue
        if uw:
            streak += 1
            if streak >= NEED_STREAK:
                start = i - NEED_STREAK + 1
                break
        else:
            streak = 0

    # Scan from end
    end = len(paths)
    streak = 0
    for j in range(len(paths) - 1, -1, -1):
        uw = is_underwater(paths[j])
        if uw is None:
            continue
        if uw:
            streak += 1
            if streak >= NEED_STREAK:
                end = j + NEED_STREAK  # exclusive
                break
        else:
            streak = 0

    if end <= start:
        # Sanity: never delete everything.
        start = 0
        end = len(paths)

    removed_head = start
    removed_tail = len(paths) - end

    if not dry_run:
        for p in paths[:start]:
            try:
                p.unlink()
            except OSError:
                pass
        for p in paths[end:]:
            try:
                p.unlink()
            except OSError:
                pass

    return (removed_head, removed_tail, end - start)


# -----------------------------------------------------------------------------
# Safety: is this segment currently being touched?
# -----------------------------------------------------------------------------

def has_ffmpeg_running_on(frames_dir: Path) -> bool:
    """Check if any ffmpeg process is writing into frames_dir."""
    try:
        r = subprocess.run(
            ["pgrep", "-af", "ffmpeg"], capture_output=True, text=True, timeout=5
        )
        for line in r.stdout.splitlines():
            if str(frames_dir) in line:
                return True
    except Exception:
        pass
    return False


def has_inference_running_on(frames_dir: Path) -> bool:
    """Check if any 05_inference.py process is running on frames_dir."""
    try:
        r = subprocess.run(
            ["pgrep", "-af", "05_inference"], capture_output=True, text=True, timeout=5
        )
        for line in r.stdout.splitlines():
            if str(frames_dir) in line:
                return True
    except Exception:
        pass
    return False


def has_pending_inference_job(conn, mission_id: int, auv_id: str, segment: str) -> bool:
    """Check state.db for queued/running 05_inference job on this segment."""
    row = conn.execute(
        "SELECT status FROM jobs WHERE mission_id=? AND auv_id=? "
        "AND segment_label=? AND stage='05_inference'",
        (mission_id, auv_id, segment),
    ).fetchone()
    if row is None:
        return False
    return row["status"] in ("queued", "running")


def frame_count_is_stable(frames_dir: Path, wait_s: float = 5.0) -> bool:
    """Return True if the frame count doesn't change over wait_s."""
    n1 = sum(1 for _ in frames_dir.glob("frame_*.jpg"))
    time.sleep(wait_s)
    n2 = sum(1 for _ in frames_dir.glob("frame_*.jpg"))
    return n1 == n2


# -----------------------------------------------------------------------------
# QC re-scoring (mirrors stage 04 qc_segment)
# -----------------------------------------------------------------------------

def qc_segment(frames_dir: Path, sample_rate: int = QC_SAMPLE_RATE) -> dict | None:
    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"  [04b] qc.json write failed: {e}")
    return agg


# -----------------------------------------------------------------------------
# Main per-segment driver
# -----------------------------------------------------------------------------

def process_segment(mission_name: str, auv_id: str, segment: str,
                    frames_dir: Path, dry_run: bool, conn) -> dict:
    result = {
        "auv_id": auv_id,
        "segment": segment,
        "frames_dir": str(frames_dir),
        "skipped": False,
        "head_removed": 0,
        "tail_removed": 0,
        "remaining": 0,
        "before_total": 0,
        "before_bottom_pct": None,
        "after_bottom_pct": None,
        "status_before": None,
        "status_after": None,
    }

    if not frames_dir.is_dir():
        result["skipped"] = True
        result["reason"] = "no_frames_dir"
        return result

    # Safety checks
    if has_ffmpeg_running_on(frames_dir):
        result["skipped"] = True
        result["reason"] = "ffmpeg_running"
        print(f"  [04b] SKIP {auv_id}/{segment}: ffmpeg still extracting")
        return result

    if has_inference_running_on(frames_dir):
        result["skipped"] = True
        result["reason"] = "inference_running_proc"
        print(f"  [04b] SKIP {auv_id}/{segment}: 05_inference process running")
        return result

    # Look up mission_id + current 04 job
    mission_row = conn.execute(
        "SELECT id FROM missions WHERE name=?", (mission_name,)
    ).fetchone()
    if not mission_row:
        result["skipped"] = True
        result["reason"] = "mission_not_in_db"
        return result
    mission_id = mission_row["id"]

    if has_pending_inference_job(conn, mission_id, auv_id, segment):
        result["skipped"] = True
        result["reason"] = "inference_job_pending"
        print(f"  [04b] SKIP {auv_id}/{segment}: 05_inference queued/running in DB")
        return result

    if not frame_count_is_stable(frames_dir, wait_s=5.0):
        result["skipped"] = True
        result["reason"] = "frame_count_unstable"
        print(f"  [04b] SKIP {auv_id}/{segment}: frame count not stable")
        return result

    # Snapshot before
    before_paths = sorted(frames_dir.glob("frame_*.jpg"))
    result["before_total"] = len(before_paths)
    job04_row = conn.execute(
        "SELECT id, status FROM jobs WHERE mission_id=? AND auv_id=? "
        "AND segment_label=? AND stage='04_frame_extract'",
        (mission_id, auv_id, segment),
    ).fetchone()
    if job04_row is None:
        result["skipped"] = True
        result["reason"] = "no_04_job_in_db"
        print(f"  [04b] SKIP {auv_id}/{segment}: no 04 job row")
        return result
    result["status_before"] = job04_row["status"]

    # Read current QC if available
    qc_path = frames_dir / "qc.json"
    if qc_path.exists():
        try:
            result["before_bottom_pct"] = json.loads(qc_path.read_text()).get("bottom_visible_pct")
        except Exception:
            pass

    # Trim
    head, tail, remaining = trim_segment(frames_dir, dry_run=dry_run)
    result["head_removed"] = head
    result["tail_removed"] = tail
    result["remaining"] = remaining

    # Re-QC if not dry-run and something was trimmed (or always to keep metrics fresh)
    after_agg = None
    if not dry_run and (head > 0 or tail > 0):
        after_agg = qc_segment(frames_dir)
    elif dry_run:
        # In dry-run, don't touch qc.json; compute aggregate from remaining slice in-memory
        remaining_paths = sorted(frames_dir.glob("frame_*.jpg"))[head: len(before_paths) - tail]
        sampled = remaining_paths[::max(1, QC_SAMPLE_RATE)]
        per_frame = [s for s in (score_image_file(f) for f in sampled) if s is not None]
        if per_frame:
            after_agg = qc_aggregate(per_frame)

    if after_agg is not None:
        result["after_bottom_pct"] = after_agg.get("bottom_visible_pct")

    if dry_run:
        print(
            f"  [04b] DRY {auv_id}/{segment}: head={head} tail={tail} "
            f"remaining={remaining} (before={len(before_paths)}, "
            f"bottom_pct {result['before_bottom_pct']}→{result['after_bottom_pct']})"
        )
        return result

    # Update DB: job row + metrics
    job_id = job04_row["id"]
    bottom_pct = after_agg.get("bottom_visible_pct") if after_agg else None

    if bottom_pct is not None and bottom_pct >= QC_BOTTOM_OK_PCT:
        new_status = "done"
        err_msg = None
    elif bottom_pct is not None:
        new_status = "degraded"
        err_msg = f"bottom_visible_pct={bottom_pct}% <{QC_BOTTOM_OK_PCT}% (after trim)"
    else:
        new_status = job04_row["status"]
        err_msg = None

    upsert_job(
        conn, mission_id, auv_id, segment, "04_frame_extract",
        status=new_status,
        output_path=str(frames_dir),
        error_msg=err_msg,
    )
    record_metric(conn, job_id, "trimmed_head", value=float(head))
    record_metric(conn, job_id, "trimmed_tail", value=float(tail))
    record_metric(conn, job_id, "frames_after_trim", value=float(remaining))
    if after_agg:
        for k in (
            "frames_total", "frames_bottom_visible", "frames_out_of_water",
            "frames_turbid", "frames_water_no_bottom",
        ):
            if k in after_agg:
                record_metric(conn, job_id, k, value=float(after_agg[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",
            )

    result["status_after"] = new_status
    print(
        f"  [04b] {auv_id}/{segment}: trimmed head={head} tail={tail} "
        f"remaining={remaining}, bottom_pct={bottom_pct}% ({result['status_before']}→{new_status})"
    )
    return result


# -----------------------------------------------------------------------------
# Discovery + CLI
# -----------------------------------------------------------------------------

def find_segments(mission_name: str, auv_filter: str | None,
                  segment_filter: str | None) -> list[tuple[str, str, Path]]:
    base = PIPELINE_BASE / "data" / mission_name / "frames"
    out: list[tuple[str, str, Path]] = []
    if not base.is_dir():
        return out
    for auv_dir in sorted(base.iterdir()):
        if not auv_dir.is_dir():
            continue
        if auv_filter and auv_dir.name != auv_filter:
            continue
        for seg_dir in sorted(auv_dir.iterdir()):
            if not seg_dir.is_dir():
                continue
            if segment_filter and seg_dir.name != segment_filter:
                continue
            out.append((auv_dir.name, seg_dir.name, seg_dir))
    return out


def main():
    ap = argparse.ArgumentParser(description="Stage 04b — Trim hors-eau head/tail frames")
    ap.add_argument("--mission", default="20260505-Lepradet")
    ap.add_argument("--auv")
    ap.add_argument("--segment")
    ap.add_argument("--dry-run", action="store_true")
    args = ap.parse_args()

    init_db()

    segments = find_segments(args.mission, args.auv, args.segment)
    if not segments:
        print(f"[04b] No segments found under {args.mission}")
        sys.exit(1)

    print(f"[04b] Mission={args.mission}  segments={len(segments)}  dry_run={args.dry_run}")
    results: list[dict] = []
    with get_conn() as conn:
        for auv_id, segment, frames_dir in segments:
            try:
                r = process_segment(args.mission, auv_id, segment, frames_dir,
                                    args.dry_run, conn)
            except Exception as e:
                r = {"auv_id": auv_id, "segment": segment, "error": str(e),
                     "skipped": True}
                print(f"  [04b] ERR {auv_id}/{segment}: {e}")
            results.append(r)

    # Summary
    print("\n=== Stage 04b summary ===")
    upgraded = [r for r in results
                if r.get("status_before") == "degraded" and r.get("status_after") == "done"]
    still_degraded = [r for r in results
                      if r.get("status_after") == "degraded"]
    skipped = [r for r in results if r.get("skipped")]
    print(f"Upgraded degraded→done : {len(upgraded)}")
    for r in upgraded:
        print(f"  + {r['auv_id']}/{r['segment']} "
              f"({r['before_bottom_pct']}%→{r['after_bottom_pct']}%, "
              f"trim head={r['head_removed']} tail={r['tail_removed']})")
    print(f"Still degraded         : {len(still_degraded)}")
    print(f"Skipped                : {len(skipped)}")
    for r in skipped:
        print(f"  - {r['auv_id']}/{r['segment']}: {r.get('reason', 'unknown')}")


if __name__ == "__main__":
    main()
