#!/usr/bin/env python3
"""Stage 01 — Ingest mission folder from SSD.

Scans /mnt/ssd/<mission>/raw_data/ and builds a manifest:
- Videos per AUV+GoPro segment (from medias/videos/)
- USBL/bag sessions per AUV (from logs/SUB/bag/*.mcap)
- Auto-detects AUV ID mapping (AUV010↔AUV210 etc.) by timestamp proximity

Usage:
    python3 01_ingest.py /mnt/ssd/20260505-Lepradet --name 20260505-Lepradet
    python3 01_ingest.py /mnt/ssd/20260505-Lepradet --name 20260505-Lepradet --out /home/cosma/cosma-pipeline
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import subprocess
import sys
from datetime import datetime, timedelta
from pathlib import Path

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

LOG_DIR = Path(os.environ.get("COSMA_PIPELINE_LOGS", "/home/cosma/cosma-pipeline/logs"))
LOG_DIR.mkdir(parents=True, exist_ok=True)

# AUV ID normalization: GP folder uses AUV2xx, bag files use AUV0xx
# e.g. AUV210 <-> AUV010, AUV213 <-> AUV013, AUV212 <-> AUV012
AUV_FOLDER_RE = re.compile(r"GP(\d+)[_-]AUV(\d+)", re.I)
AUV_BAG_RE    = re.compile(r"auv(\d+)", re.I)
AUV_CSV_RE    = re.compile(r"AUV(\d+)", re.I)


def normalize_auv(raw_id: str) -> str:
    """Normalize AUV IDs: strip leading zeros or map 0xx -> 2xx heuristic.
    Returns canonical form like AUV010, AUV013, AUV210, AUV212 as-is.
    We keep original for now and detect mapping via timestamp cross-correlation.
    """
    m = re.search(r"\d+", raw_id)
    if not m:
        return raw_id.upper()
    n = int(m.group())
    return f"AUV{n:03d}"


def exif_create_date(path: Path) -> datetime | None:
    try:
        out = subprocess.check_output(
            ["exiftool", "-s3", "-CreateDate", "-api", "QuickTimeUTC=1", str(path)],
            stderr=subprocess.DEVNULL, text=True, timeout=10
        ).strip()
        if not out:
            return None
        out = re.sub(r'[+-]\d{2}:\d{2}$', '', out).strip()
        return datetime.strptime(out, "%Y:%m:%d %H:%M:%S")
    except Exception:
        return None


def exif_duration_s(path: Path) -> float | None:
    try:
        out = subprocess.check_output(
            ["exiftool", "-s3", "-Duration#", str(path)],
            stderr=subprocess.DEVNULL, text=True, timeout=10
        ).strip()
        return float(out) if out else None
    except Exception:
        return None


def scan_videos(raw_data: Path) -> dict[str, list[dict]]:
    """Scan medias/videos/ and return dict {auv_id: [video_info, ...]}.
    Handles both GP1-AUV210 and GP1_AUV210 naming conventions.
    """
    videos_dir = raw_data / "medias" / "videos"
    if not videos_dir.exists():
        return {}

    result: dict[str, list[dict]] = {}
    for folder in sorted(videos_dir.iterdir()):
        if not folder.is_dir():
            continue
        m = AUV_FOLDER_RE.search(folder.name)
        if not m:
            continue
        gopro_n = int(m.group(1))
        auv_id = normalize_auv(m.group(2))

        mp4_files = sorted(folder.glob("*.MP4")) + sorted(folder.glob("*.mp4"))
        for mp4 in mp4_files:
            create_date = exif_create_date(mp4)
            duration = exif_duration_s(mp4)
            info = {
                "path": str(mp4),
                "gopro": gopro_n,
                "auv_id": auv_id,
                "filename": mp4.name,
                "create_date": create_date.isoformat() if create_date else None,
                "duration_s": duration,
                "size_mb": round(mp4.stat().st_size / 1e6, 1),
            }
            result.setdefault(auv_id, []).append(info)

    return result


def scan_bags(raw_data: Path) -> dict[str, list[dict]]:
    """Scan logs/SUB/bag/ for MCAP files grouped by session+AUV."""
    bag_dir = raw_data / "logs" / "SUB" / "bag"
    if not bag_dir.exists():
        return {}

    result: dict[str, list[dict]] = {}
    for session_dir in sorted(bag_dir.iterdir()):
        if not session_dir.is_dir():
            continue
        # dir name: 20260505_074718_AUV013
        m = re.match(r"(\d{8}_\d{6})_AUV(\d+)", session_dir.name)
        if not m:
            continue
        ts_str = m.group(1)
        auv_id = normalize_auv(m.group(2))
        try:
            ts = datetime.strptime(ts_str, "%Y%m%d_%H%M%S")
        except ValueError:
            ts = None

        mcap_files = sorted(session_dir.glob("*.mcap"))
        total_size = sum(f.stat().st_size for f in mcap_files)
        non_empty = [str(f) for f in mcap_files if f.stat().st_size > 0]

        if not non_empty:
            continue

        session_info = {
            "session": session_dir.name,
            "auv_id": auv_id,
            "timestamp": ts.isoformat() if ts else None,
            "mcap_files": non_empty,
            "total_mb": round(total_size / 1e6, 1),
        }
        result.setdefault(auv_id, []).append(session_info)

    return result


def scan_usbl_csv(raw_data: Path) -> dict[str, list[dict]]:
    """Scan logs/SUB/log/ for *_usbl.csv files.
    Note: these are raw serial byte logs, not lat/lon CSV.
    We record them for reference; actual nav comes from MCAP bags.
    """
    log_dir = raw_data / "logs" / "SUB" / "log"
    if not log_dir.exists():
        return {}

    result: dict[str, list[dict]] = {}
    for f in sorted(log_dir.glob("*_usbl.csv")):
        m = AUV_CSV_RE.search(f.name)
        if not m:
            continue
        auv_id = normalize_auv(m.group(1))
        # parse timestamp from filename: 2026-05-05_08-16-00_AUV010_usbl.csv
        ts_m = re.match(r"(\d{4}-\d{2}-\d{2}_\d{2}-\d{2}-\d{2})", f.name)
        ts = None
        if ts_m:
            try:
                ts = datetime.strptime(ts_m.group(1), "%Y-%m-%d_%H-%M-%S")
            except ValueError:
                pass
        result.setdefault(auv_id, []).append({
            "path": str(f),
            "auv_id": auv_id,
            "timestamp": ts.isoformat() if ts else None,
            "size_kb": round(f.stat().st_size / 1e3, 1),
            "format": "raw_serial",  # NOT lat/lon CSV
        })

    return result


def detect_auv_mapping(videos: dict, bags: dict) -> dict[str, str]:
    """Detect AUV ID mapping between video folders (AUV2xx) and bag sessions (AUV0xx).

    Heuristic: if video AUV2xx and bag AUV0xx share same last 2 digits and
    timestamps are within 60s → they are the same physical AUV.

    Returns: {video_auv_id: bag_auv_id} e.g. {"AUV210": "AUV010"}
    """
    mapping: dict[str, str] = {}

    for vid_auv in videos:
        vid_ts_list = []
        for v in videos[vid_auv]:
            if v.get("create_date"):
                try:
                    vid_ts_list.append(datetime.fromisoformat(v["create_date"]))
                except Exception:
                    pass
        if not vid_ts_list:
            continue

        vid_digits = vid_auv[-2:]  # last 2 digits of AUV2xx

        best_bag_auv = None
        best_delta = timedelta(seconds=999)

        for bag_auv in bags:
            # check same last 2 digits
            if bag_auv[-2:] != vid_digits:
                continue
            for sess in bags[bag_auv]:
                if sess.get("timestamp"):
                    try:
                        bag_ts = datetime.fromisoformat(sess["timestamp"])
                    except Exception:
                        continue
                    for vt in vid_ts_list:
                        delta = abs(vt - bag_ts)
                        if delta < best_delta:
                            best_delta = delta
                            best_bag_auv = bag_auv

        if best_bag_auv and best_delta < timedelta(seconds=3600):
            mapping[vid_auv] = best_bag_auv

    return mapping


def group_video_segments(video_list: list[dict], gap_min: int = 5) -> list[dict]:
    """Group consecutive videos into segments by timestamp gap."""
    sorted_vids = sorted(
        [v for v in video_list if v.get("create_date")],
        key=lambda x: x["create_date"]
    )
    if not sorted_vids:
        return [{"videos": video_list, "label": "seg_unknown", "total_s": None}]

    segments = []
    current: list[dict] = [sorted_vids[0]]

    for vid in sorted_vids[1:]:
        prev = current[-1]
        try:
            prev_end = datetime.fromisoformat(prev["create_date"])
            if prev.get("duration_s"):
                prev_end += timedelta(seconds=prev["duration_s"])
            cur_start = datetime.fromisoformat(vid["create_date"])
            gap = (cur_start - prev_end).total_seconds() / 60
            if gap > gap_min:
                segments.append(_finalize_segment(current))
                current = [vid]
            else:
                current.append(vid)
        except Exception:
            current.append(vid)

    if current:
        segments.append(_finalize_segment(current))

    return segments


def _finalize_segment(videos: list[dict]) -> dict:
    label = videos[0]["create_date"][:19].replace(":", "-").replace("T", "_") if videos[0].get("create_date") else "seg_unknown"
    total_s = sum(v["duration_s"] or 0 for v in videos)
    return {
        "label": label,
        "videos": videos,
        "total_s": total_s,
        "start": videos[0].get("create_date"),
        "end": videos[-1].get("create_date"),
    }


def build_manifest(mission_path: Path, gap_min: int = 5) -> dict:
    raw_data = mission_path / "raw_data"
    if not raw_data.exists():
        # try direct
        raw_data = mission_path

    print(f"[01_ingest] scanning {raw_data} ...")

    videos = scan_videos(raw_data)
    bags   = scan_bags(raw_data)
    csvs   = scan_usbl_csv(raw_data)
    auv_map = detect_auv_mapping(videos, bags)

    # Build per-AUV segments
    auv_segments: dict[str, list[dict]] = {}
    for auv_id, vid_list in videos.items():
        segs = group_video_segments(vid_list, gap_min=gap_min)
        auv_segments[auv_id] = segs

    # Compute AUVs with real data
    auv_ids_with_video = sorted(videos.keys())
    auv_ids_with_bags  = sorted(bags.keys())

    total_video_s = sum(
        seg["total_s"] or 0
        for segs in auv_segments.values()
        for seg in segs
    )

    manifest = {
        "mission": mission_path.name,
        "ssd_path": str(mission_path),
        "generated_at": now_iso(),
        "auv_ids_video": auv_ids_with_video,
        "auv_ids_bags": auv_ids_with_bags,
        "auv_mapping": auv_map,
        "total_video_s": round(total_video_s),
        "segments_per_auv": auv_segments,
        "bag_sessions_per_auv": bags,
        "usbl_csv_per_auv": csvs,
        "notes": {
            "usbl_csv_format": "raw_serial_bytes",
            "nav_source": "mcap_bags",
        },
    }

    return manifest


def ingest(mission_path: Path, mission_name: str, out_dir: Path,
           gap_min: int = 5) -> dict:
    out_dir.mkdir(parents=True, exist_ok=True)
    manifest_path = out_dir / mission_name / "manifest.json"

    # Idempotency check
    if manifest_path.exists():
        existing = json.loads(manifest_path.read_text())
        chk = hashlib.sha256(mission_path.name.encode()).hexdigest()[:8]
        print(f"[01_ingest] manifest exists (checksum {chk}), skipping scan")
        return existing

    manifest = build_manifest(mission_path, gap_min=gap_min)

    # Save manifest
    manifest_path.parent.mkdir(parents=True, exist_ok=True)
    manifest_path.write_text(json.dumps(manifest, indent=2))
    print(f"[01_ingest] manifest saved: {manifest_path}")

    # Write to DB
    init_db()
    with get_conn() as conn:
        upsert_mission(conn, mission_name, str(mission_path),
                       status="ingested", manifest=json.dumps(manifest))

    return manifest


def main():
    ap = argparse.ArgumentParser(description="Stage 01 — Ingest mission from SSD")
    ap.add_argument("mission_path", type=Path, help="Path to mission folder (e.g. /mnt/ssd/20260505-Lepradet)")
    ap.add_argument("--name", type=str, default=None, help="Mission name (defaults to folder name)")
    ap.add_argument("--out", type=Path, default=Path("/home/cosma/cosma-pipeline"), help="Output base dir")
    ap.add_argument("--gap-min", type=int, default=5, help="Gap in minutes to split video segments")
    args = ap.parse_args()

    mission_name = args.name or args.mission_path.name
    manifest = ingest(args.mission_path, mission_name, args.out, args.gap_min)

    print(f"\n=== Ingest summary for {mission_name} ===")
    print(f"AUVs with video: {manifest['auv_ids_video']}")
    print(f"AUVs with bags:  {manifest['auv_ids_bags']}")
    print(f"AUV mapping:     {manifest['auv_mapping']}")
    print(f"Total video:     {manifest['total_video_s']}s")
    print(f"Segments:")
    for auv, segs in manifest["segments_per_auv"].items():
        for seg in segs:
            n_vids = len(seg["videos"])
            dur = f"{seg['total_s']:.0f}s" if seg["total_s"] else "?"
            print(f"  {auv} / {seg['label']}  {n_vids} videos  {dur}")


if __name__ == "__main__":
    main()
