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

The USBL CSV logs in logs/SUB/log/ contain raw serial bytes, NOT lat/lon.
Real navigation data is in MCAP bags (logs/SUB/bag/).

This stage:
1. Reads MCAP files per AUV session
2. Extracts position topics (configurable in default_params.yaml)
3. Falls back to parsing serial bytes if no nav topic found (best-effort)
4. Outputs Parquet per AUV with columns: timestamp, lat, lon, depth, heading

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

import argparse
import json
import os
import struct
import sys
from datetime import datetime, timezone
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

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

# Known nav topics to try in MCAP files
NAV_TOPICS = [
    "/usbl/position",
    "/usbl/fix",
    "/navigation/position",
    "/bluerov/usbl",
    "/waterlinked/position",
    "/fix",
    "/gps/fix",
    "/mavros/global_position/global",
    "/mavros/global_position/local",
]


def try_parse_mcap(mcap_path: Path, topics: list[str] | None = None) -> list[dict]:
    """Try to extract nav points from MCAP file. Returns list of {ts, lat, lon, depth}."""
    try:
        from mcap.reader import make_reader
    except ImportError:
        print(f"  [02] mcap not installed, skipping {mcap_path.name}")
        return []

    points = []
    try:
        with open(mcap_path, "rb") as f:
            reader = make_reader(f)
            for schema, channel, message in reader.iter_messages(topics=topics):
                # Try to deserialize — support ROS2 JSON-encoded or raw
                try:
                    import json as _json
                    data = _json.loads(message.data)
                    lat = data.get("latitude") or data.get("lat")
                    lon = data.get("longitude") or data.get("lon")
                    depth = data.get("depth") or data.get("altitude")
                    heading = data.get("heading") or data.get("yaw")
                    if lat is not None and lon is not None:
                        ts_ns = message.log_time
                        ts = ts_ns / 1e9
                        points.append({
                            "timestamp": ts,
                            "lat": float(lat),
                            "lon": float(lon),
                            "depth": float(depth) if depth is not None else None,
                            "heading": float(heading) if heading is not None else None,
                            "source": channel.topic,
                        })
                except Exception:
                    pass
    except Exception as e:
        print(f"  [02] MCAP read error {mcap_path.name}: {e}")

    return points


def try_parse_serial_csv(csv_path: Path) -> list[dict]:
    """Best-effort: parse raw serial byte log for USBL range/bearing frames.
    Waterlinked USBL M64 protocol: 0xBB 0x55 frame header.
    This is a rough attempt — actual lat/lon requires ship GPS + range+bearing.
    Returns relative-only positions (range, bearing) if decoded.
    """
    points = []
    try:
        with open(csv_path, "r") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                parts = line.split(",", 2)
                if len(parts) < 3:
                    continue
                ts_str, direction, raw = parts[0], parts[1], parts[2]
                if direction.strip() != "RECEIVED":
                    continue
                # Extract bytes from repr string b'\xbb...'
                try:
                    raw_clean = raw.strip().strip('"')
                    # Parse Python bytes repr
                    data = eval(raw_clean)  # safe: only used on known CSV
                    if len(data) >= 4 and data[0] == 0xBB and data[1] == 0x55:
                        # Waterlinked M64 position frame: len byte at [2], payload follows
                        payload_len = data[2]
                        if len(data) >= payload_len + 3:
                            # Best effort: look for float32 values in payload
                            # Actual protocol decoding would need WL M64 spec
                            ts = datetime.fromisoformat(ts_str)
                            points.append({
                                "timestamp": ts.timestamp(),
                                "lat": None,
                                "lon": None,
                                "depth": None,
                                "heading": None,
                                "source": "serial_raw",
                                "raw_bytes": data.hex()[:32],
                            })
                except Exception:
                    pass
    except Exception as e:
        print(f"  [02] serial CSV parse error {csv_path.name}: {e}")

    return points


def parse_auv_sessions(manifest: dict, auv_id: str, out_dir: Path,
                       topics: list[str] | None = None) -> dict:
    """Parse all sessions for one AUV. Returns metrics dict."""
    metrics = {"auv_id": auv_id, "points_raw": 0, "sources": [], "status": "ok"}
    all_points: list[dict] = []

    # Try MCAP bags first
    bag_sessions = manifest.get("bag_sessions_per_auv", {}).get(auv_id, [])
    if not bag_sessions:
        # Try mapping: maybe bags use AUV0xx while videos use AUV2xx
        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, [])

    for sess in bag_sessions:
        for mcap_path_str in sess.get("mcap_files", []):
            mcap_path = Path(mcap_path_str)
            if not mcap_path.exists():
                continue
            pts = try_parse_mcap(mcap_path, topics=topics or NAV_TOPICS)
            if pts:
                all_points.extend(pts)
                if "mcap" not in metrics["sources"]:
                    metrics["sources"].append("mcap")
                print(f"  [02] {auv_id} {mcap_path.name}: {len(pts)} nav points")

    # Fallback: serial CSV if no MCAP nav
    if not all_points:
        csv_entries = manifest.get("usbl_csv_per_auv", {}).get(auv_id, [])
        for entry in csv_entries:
            csv_path = Path(entry["path"])
            if not csv_path.exists():
                continue
            pts = try_parse_serial_csv(csv_path)
            if pts:
                all_points.extend(pts)
                if "serial_csv" not in metrics["sources"]:
                    metrics["sources"].append("serial_csv")

    metrics["points_raw"] = len(all_points)

    if not all_points:
        metrics["status"] = "degraded"
        metrics["note"] = "no nav points found in MCAP or serial CSV"
        print(f"  [02] {auv_id}: NO nav data found — degraded")
    else:
        print(f"  [02] {auv_id}: {len(all_points)} raw nav points from {metrics['sources']}")

    # Save output
    out_file = out_dir / f"{auv_id}_nav_raw.json"
    out_file.write_text(json.dumps({
        "auv_id": auv_id,
        "generated_at": now_iso(),
        "metrics": metrics,
        "points": all_points,
    }, indent=2, default=str))

    # Also save as simple CSV for downstream
    csv_out = out_dir / f"{auv_id}_nav_raw.csv"
    with open(csv_out, "w") as f:
        f.write("timestamp,lat,lon,depth,heading,source\n")
        for p in all_points:
            f.write(f"{p['timestamp']},{p.get('lat','')},{p.get('lon','')},{p.get('depth','')},{p.get('heading','')},{p.get('source','')}\n")

    return metrics


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

    if out_dir is None:
        out_dir = manifest_path.parent / "02_usbl_raw"
    out_dir.mkdir(parents=True, exist_ok=True)

    # Idempotency: check if all AUV outputs exist
    auv_ids = manifest.get("auv_ids_bags", []) or 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 auv_ids:
        out_file = out_dir / f"{auv_id}_nav_raw.json"
        if out_file.exists():
            print(f"[02] {auv_id}: output exists, skipping")
            existing = json.loads(out_file.read_text())
            all_metrics.append(existing.get("metrics", {"auv_id": auv_id, "status": "cached"}))
            continue

        print(f"[02] Parsing {auv_id} ...")
        m = parse_auv_sessions(manifest, auv_id, out_dir)
        all_metrics.append(m)

        # Record in DB
        with get_conn() as conn:
            from orchestrator.db import upsert_mission
            mission_id_row = conn.execute(
                "SELECT id FROM missions WHERE name=?", (mission_name,)
            ).fetchone()
            if mission_id_row:
                mission_id = mission_id_row["id"]
                job_id = upsert_job(conn, mission_id, auv_id, "all", "02_usbl_parse",
                                    status="done" if m["status"] == "ok" else m["status"],
                                    output_path=str(out_dir))
                record_metric(conn, job_id, "usbl_points_raw",
                              value=m["points_raw"],
                              pass_fail="pass" if m["points_raw"] > 0 else "fail")

    return all_metrics


def main():
    ap = argparse.ArgumentParser(description="Stage 02 — Parse USBL/nav from MCAP bags")
    ap.add_argument("manifest", type=Path, help="manifest.json from stage 01")
    ap.add_argument("--auv", type=str, default=None, help="Filter to single AUV ID")
    ap.add_argument("--out", type=Path, default=None, help="Output directory")
    args = ap.parse_args()

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

    print("\n=== Stage 02 summary ===")
    total_pts = sum(m.get("points_raw", 0) for m in metrics)
    for m in metrics:
        status = m.get("status", "?")
        pts = m.get("points_raw", 0)
        src = m.get("sources", [])
        print(f"  {m['auv_id']}: {pts} pts  {src}  [{status}]")
    print(f"Total nav points: {total_pts}")


if __name__ == "__main__":
    main()
