#!/usr/bin/env python3
"""Stage 02 — Parse navigation from ROS2 MCAP bag files.

Extracts per-AUV trajectories from MCAP bags using mcap_ros2:
- /mavros/global_position/global  → NavSatFix (lat, lon, alt)
- /mavros/imu/data               → Imu (qx, qy, qz, qw)
- /mavros/imu/static_pressure    → FluidPressure (pressure_pa)

Joins on nearest timestamp (tolerance 100ms).
Saves parquet: ~/cosma-pipeline/data/<mission>/nav/<AUV>_<segment>.parquet

Fallback: if no MCAP GPS data, marks as degraded=True (GPS=0 under water is normal).

Usage:
    python3 02_nav_parse.py /home/cosma/cosma-pipeline/20260505-Lepradet/manifest.json
    python3 02_nav_parse.py /path/manifest.json --auv AUV013
"""
from __future__ import annotations

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

import numpy as np

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"))
NAV_TOPICS = [
    "/mavros/global_position/global",
    "/mavros/imu/data",
    "/mavros/imu/static_pressure",
]


def parse_mcap_segment(mcap_files: list[Path]) -> dict[str, list]:
    """Extract raw topic data from a list of MCAP files (one session/segment).
    Returns dict keyed by topic -> list of (ts_ns, data_dict).
    """
    from mcap_ros2.reader import read_ros2_messages

    topic_data: dict[str, list] = {t: [] for t in NAV_TOPICS}

    for mcap_path in mcap_files:
        if not mcap_path.exists():
            continue
        try:
            for msg in read_ros2_messages(str(mcap_path), topics=NAV_TOPICS):
                topic = msg.channel.topic
                m = msg.ros_msg
                ts_ns = int(msg.log_time.timestamp() * 1e9)

                if topic == "/mavros/global_position/global":
                    topic_data[topic].append((ts_ns, {
                        "lat": float(m.latitude),
                        "lon": float(m.longitude),
                        "alt": float(m.altitude),
                    }))
                elif topic == "/mavros/imu/data":
                    topic_data[topic].append((ts_ns, {
                        "qx": float(m.orientation.x),
                        "qy": float(m.orientation.y),
                        "qz": float(m.orientation.z),
                        "qw": float(m.orientation.w),
                    }))
                elif topic == "/mavros/imu/static_pressure":
                    topic_data[topic].append((ts_ns, {
                        "pressure_pa": float(m.fluid_pressure),
                    }))
        except Exception as e:
            print(f"  [02] Error reading {mcap_path.name}: {e}")

    return topic_data


def join_topics(topic_data: dict[str, list], tol_ns: int = 100_000_000) -> list[dict]:
    """Join NavSatFix + Imu + FluidPressure on nearest timestamp (100ms tol).
    Base timeline = NavSatFix if available, else Imu.
    """
    import pandas as pd

    nav_pts = topic_data.get("/mavros/global_position/global", [])
    imu_pts = topic_data.get("/mavros/imu/data", [])
    pres_pts = topic_data.get("/mavros/imu/static_pressure", [])

    if not nav_pts and not imu_pts:
        return []

    # Build DataFrames
    if nav_pts:
        df_nav = pd.DataFrame([{"ts_ns": ts, **d} for ts, d in nav_pts])
    else:
        df_nav = pd.DataFrame(columns=["ts_ns", "lat", "lon", "alt"])

    if imu_pts:
        df_imu = pd.DataFrame([{"ts_ns": ts, **d} for ts, d in imu_pts])
    else:
        df_imu = pd.DataFrame(columns=["ts_ns", "qx", "qy", "qz", "qw"])

    if pres_pts:
        df_pres = pd.DataFrame([{"ts_ns": ts, **d} for ts, d in pres_pts])
    else:
        df_pres = pd.DataFrame(columns=["ts_ns", "pressure_pa"])

    # Use nav as base if it has data, else imu
    base_df = df_nav if len(df_nav) > 0 else df_imu
    base_df = base_df.sort_values("ts_ns").reset_index(drop=True)

    # Merge-as-of for IMU
    result = base_df.copy()
    if len(df_imu) > 0:
        df_imu_s = df_imu.sort_values("ts_ns").reset_index(drop=True)
        # Simple nearest-neighbor join
        imu_ts = df_imu_s["ts_ns"].values
        for col in ["qx", "qy", "qz", "qw"]:
            result[col] = np.nan
        for i, row_ts in enumerate(result["ts_ns"].values):
            idx = np.argmin(np.abs(imu_ts - row_ts))
            if abs(imu_ts[idx] - row_ts) <= tol_ns:
                for col in ["qx", "qy", "qz", "qw"]:
                    result.at[i, col] = float(df_imu_s.at[idx, col])

    # Merge pressure
    if len(df_pres) > 0:
        df_pres_s = df_pres.sort_values("ts_ns").reset_index(drop=True)
        pres_ts = df_pres_s["ts_ns"].values
        result["pressure_pa"] = np.nan
        for i, row_ts in enumerate(result["ts_ns"].values):
            idx = np.argmin(np.abs(pres_ts - row_ts))
            if abs(pres_ts[idx] - row_ts) <= tol_ns:
                result.at[i, "pressure_pa"] = float(df_pres_s.at[idx, "pressure_pa"])

    # Ensure all columns exist
    for col in ["lat", "lon", "alt", "qx", "qy", "qz", "qw", "pressure_pa"]:
        if col not in result.columns:
            result[col] = np.nan

    return result.to_dict("records")


def parse_auv(manifest: dict, auv_id: str, out_dir: Path) -> dict:
    """Parse all MCAP sessions for one AUV. Returns metrics."""
    from pathlib import Path as P

    metrics = {
        "auv_id": auv_id,
        "segments": [],
        "total_points": 0,
        "degraded": False,
        "status": "ok",
    }

    bag_sessions = manifest.get("bag_sessions_per_auv", {}).get(auv_id, [])
    if not bag_sessions:
        auv_map = manifest.get("auv_mapping", {})
        bag_auv = auv_map.get(auv_id)
        if bag_auv:
            bag_sessions = manifest.get("bag_sessions_per_auv", {}).get(bag_auv, [])

    if not bag_sessions:
        # Build from raw SSD structure
        ssd_path = P(manifest.get("ssd_path", "/mnt/ssd") + "/" + manifest["mission"].split("-")[0] + "-" + manifest["mission"].split("-")[1] if "-" in manifest["mission"] else manifest.get("ssd_path", "/mnt/ssd"))
        auv_num = auv_id.replace("AUV", "0")  # AUV013 -> 0013? No: AUV013 -> AUV013
        bag_root = P(manifest.get("ssd_path", "/mnt/ssd")) / "raw_data/logs/SUB/bag"
        sessions = sorted(bag_root.glob(f"*_{auv_id}"))
        bag_sessions = [{"label": s.name, "mcap_files": [str(f) for f in sorted(s.glob("*.mcap"))]} for s in sessions]

    import pandas as pd

    all_points_total = 0
    for sess in bag_sessions:
        label = sess.get("session", sess.get("label", "unknown"))
        mcap_files = [P(f) for f in sess.get("mcap_files", [])]
        if not mcap_files:
            continue

        out_parquet = out_dir / f"{auv_id}_{label}.parquet"
        if out_parquet.exists():
            df_ex = pd.read_parquet(out_parquet)
            n = len(df_ex)
            print(f"  [02] {auv_id}/{label}: cached ({n} pts)")
            all_points_total += n
            metrics["segments"].append({"label": label, "points": n, "cached": True})
            continue

        print(f"  [02] {auv_id}/{label}: parsing {len(mcap_files)} MCAP files...")
        topic_data = parse_mcap_segment(mcap_files)
        points = join_topics(topic_data)

        if not points:
            print(f"  [02] {auv_id}/{label}: no data")
            metrics["segments"].append({"label": label, "points": 0, "degraded": True})
            metrics["degraded"] = True
            continue

        df = pd.DataFrame(points)
        n = len(df)
        # Check GPS quality
        has_gps = df["lat"].notna().any() and (df["lat"] != 0).any()
        if not has_gps:
            print(f"  [02] {auv_id}/{label}: {n} pts, GPS=0 (degraded — AUV underwater)")
            metrics["degraded"] = True
        else:
            print(f"  [02] {auv_id}/{label}: {n} pts, GPS OK")

        df.to_parquet(out_parquet, index=False)
        all_points_total += n
        metrics["segments"].append({"label": label, "points": n, "degraded": not has_gps})

    metrics["total_points"] = all_points_total
    if all_points_total == 0:
        metrics["status"] = "degraded"
    return metrics


def parse_mission(manifest_path: Path, auv_filter: str | None = None) -> list[dict]:
    manifest = json.loads(manifest_path.read_text())
    mission_name = manifest["mission"]

    out_dir = PIPELINE_BASE / "data" / mission_name / "nav"
    out_dir.mkdir(parents=True, exist_ok=True)

    auv_ids = list(set(
        manifest.get("auv_ids_bags", []) +
        list(manifest.get("auv_mapping", {}).keys())
    ))
    if not auv_ids:
        auv_ids = manifest.get("auv_ids_video", [])
    if auv_filter:
        auv_ids = [a for a in auv_ids if a == auv_filter]

    all_metrics = []
    init_db()

    for auv_id in sorted(auv_ids):
        print(f"[02] === {auv_id} ===")
        m = parse_auv(manifest, auv_id, out_dir)
        all_metrics.append(m)

        with get_conn() as conn:
            mission_row = conn.execute("SELECT id FROM missions WHERE name=?", (mission_name,)).fetchone()
            if mission_row:
                job_id = upsert_job(conn, mission_row["id"], auv_id, "all", "02_nav_parse",
                                    status="done" if m["status"] == "ok" else m["status"],
                                    output_path=str(out_dir))
                record_metric(conn, job_id, "nav_points_total", value=m["total_points"],
                              pass_fail="pass" if m["total_points"] > 0 else "warn")

    return all_metrics


def main():
    ap = argparse.ArgumentParser(description="Stage 02 — Parse nav from MCAP bags")
    ap.add_argument("manifest", type=Path)
    ap.add_argument("--auv", type=str, default=None)
    args = ap.parse_args()

    metrics = parse_mission(args.manifest, auv_filter=args.auv)

    print("\n=== Stage 02 summary ===")
    for m in metrics:
        segs = m.get("segments", [])
        total = m.get("total_points", 0)
        deg = "DEGRADED" if m.get("degraded") else "OK"
        print(f"  {m['auv_id']}: {total} pts across {len(segs)} segments [{deg}]")


if __name__ == "__main__":
    main()
