#!/usr/bin/env python3
"""Export audited H32/G90 static windows directly to RTK--LiDAR combined data.

This adapter is for captures where multiple static stations live inside large
DLog archives instead of one directory per station.  It uses the H32 packet
host-receive UTC ticks as the common software clock, parses G90 ``$GNGGA`` and
``$GNHPR`` from one or more V2 captures, and deliberately does not require IMU.
Raw inputs are opened read-only.
"""

from __future__ import annotations

import argparse
import bisect
import csv
import json
import shutil
import struct
import sys
import zipfile
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import BinaryIO, Iterator

import numpy as np

ROOT = Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))
sys.path.insert(0, str(ROOT / "rscap_v2"))

from build_multisensor_npz import build_combined  # noqa: E402
from h32_dlog.difop import DifopAngles, parse_difop_angles  # noqa: E402
from h32_dlog.dotnet_bin import read_dotnet_string  # noqa: E402
from h32_dlog.payload_v1 import parse_difop_payload, parse_msop_batch_payload  # noqa: E402
from capture_format_v2 import read_capture  # noqa: E402
from h32_msop import iter_h32_frames_polar_from_packets  # noqa: E402
from pipeline_common_corrected import parse_rtk_capture, write_jsonl  # noqa: E402

DOTNET_UNIX_EPOCH_TICKS = 621355968000000000
TICKS_PER_SECOND = 10_000_000
LOCAL_TZ = timezone(timedelta(hours=8))
MSOP_OBJECT = "frontlidar-msop-raw"
DIFOP_OBJECT = "frontlidar-difop-raw"


@dataclass(frozen=True)
class Window:
    station_id: str
    start_ticks: int
    end_ticks: int


def local_text_to_utc_ticks(text: str) -> int:
    value = datetime.strptime(text.strip(), "%Y-%m-%d %H:%M:%S.%f").replace(tzinfo=LOCAL_TZ)
    return int(round(value.timestamp() * TICKS_PER_SECOND)) + DOTNET_UNIX_EPOCH_TICKS


def read_exact(stream: BinaryIO, length: int) -> bytes:
    value = stream.read(length)
    if len(value) != length:
        raise EOFError(f"expected {length} bytes, got {len(value)}")
    return value


def iter_zip_dobject_payloads(path: Path) -> Iterator[tuple[str, bytes]]:
    """Sequentially read DObject records from a standard Medulla DLog ZIP."""

    with zipfile.ZipFile(path) as archive:
        candidates = [
            name for name in archive.namelist()
            if name.replace("\\", "/").endswith("dobject_recording/data.bin")
        ]
        if len(candidates) != 1:
            raise ValueError(f"{path}: expected one dobject_recording/data.bin, got {candidates}")
        with archive.open(candidates[0], "r") as stream:
            while True:
                try:
                    name = read_dotnet_string(stream)
                except EOFError:
                    break
                read_exact(stream, 8)  # outer DObject tick
                read_dotnet_string(stream)  # record id
                length = struct.unpack("<i", read_exact(stream, 4))[0]
                if length < 0 or length > 128 * 1024 * 1024:
                    raise ValueError(f"{path}: invalid DObject payload length {length}")
                yield name, read_exact(stream, length)


def load_windows(path: Path) -> list[Window]:
    grouped: dict[str, list[tuple[int, int]]] = {}
    with path.open("r", encoding="utf-8-sig", newline="") as stream:
        for row in csv.DictReader(stream):
            station = row["station_id"].strip()
            grouped.setdefault(station, []).append(
                (local_text_to_utc_ticks(row["local_start"]), local_text_to_utc_ticks(row["local_end"]))
            )
    merged: list[Window] = []
    for station, ranges in grouped.items():
        current: list[list[int]] = []
        for start, end in sorted(ranges):
            if current and start <= current[-1][1]:
                current[-1][1] = max(current[-1][1], end)
            else:
                current.append([start, end])
        merged.extend(Window(station, start, end) for start, end in current)
    merged.sort(key=lambda item: item.start_ticks)
    for previous, current in zip(merged, merged[1:]):
        if current.start_ticks <= previous.end_ticks and current.station_id != previous.station_id:
            raise ValueError(f"overlapping stations: {previous} and {current}")
    return merged


def station_lookup(windows: list[Window]):
    starts = [item.start_ticks for item in windows]

    def lookup(ticks: int) -> str | None:
        index = bisect.bisect_right(starts, ticks) - 1
        if index >= 0 and ticks <= windows[index].end_ticks:
            return windows[index].station_id
        return None

    return lookup


def save_frames(
    station: str,
    packet_items: dict[tuple[str, int], tuple[int, int, bytes]],
    export_root: Path,
    angles: DifopAngles,
    source: Path,
    *,
    frame_stride: int,
    seen_frame_keys: set[tuple[str, int, int]],
) -> int:
    if not packet_items:
        return 0
    ordered = sorted(packet_items.values(), key=lambda item: (item[0], item[1]))
    frames = iter_h32_frames_polar_from_packets(
        (item[2] for item in ordered),
        host_utc_ticks=[item[0] for item in ordered],
        frame_stride=frame_stride,
        min_frame_points=100,
        min_range_m=0.3,
        max_range_m=120.0,
        vertical_deg=angles.vertical_deg,
        horizontal_deg=angles.horizontal_deg,
    )
    frames_dir = export_root / station / "frames"
    frames_dir.mkdir(parents=True, exist_ok=True)
    written = 0
    for frame in frames:
        # Overlapping archives contain identical revolutions.  The host stamp
        # and 0.1 s device bucket make the key stable without comparing points.
        key = (station, int(round(frame.host_receive_utc_ns / 10_000_000)), int(round(frame.t_start_s * 10)))
        if key in seen_frame_keys:
            continue
        seen_frame_keys.add(key)
        device_ns = int(round(frame.t_start_s * 1_000_000_000))
        destination = frames_dir / f"h32_{frame.host_receive_utc_ns}_{device_ns}.npz"
        np.savez_compressed(
            destination,
            points_raw=np.asarray(frame.points_raw, dtype=np.float32),
            frame_counter=np.asarray([len(seen_frame_keys)], dtype=np.int32),
            point_count=np.asarray([len(frame.points_raw)], dtype=np.int32),
            unix_time_ns=np.asarray([device_ns], dtype=np.int64),
            device_time_s=np.asarray([frame.t_start_s], dtype=np.float64),
            device_time_end_s=np.asarray([frame.t_end_s], dtype=np.float64),
            host_receive_utc_ns=np.asarray([frame.host_receive_utc_ns], dtype=np.int64),
            source_file_utf8=np.frombuffer(str(source.resolve()).encode("utf-8"), dtype=np.uint8),
        )
        written += 1
    return written


def scan_dlog_sources(
    sources: list[Path],
    windows: list[Window],
    export_root: Path,
    *,
    frame_stride: int,
) -> dict[str, object]:
    lookup = station_lookup(windows)
    angles: DifopAngles | None = None
    seen_packets: dict[str, set[tuple[str, int]]] = {}
    seen_frames: set[tuple[str, int, int]] = set()
    frame_counts: dict[str, int] = {}
    source_stats: list[dict[str, object]] = []

    for source_index, source in enumerate(sources, 1):
        print(f"[dlog {source_index}/{len(sources)}] {source}", flush=True)
        packets: dict[str, dict[tuple[str, int], tuple[int, int, bytes]]] = {}
        msop_batches = difop_records = selected_packets = duplicates = 0
        for object_name, payload in iter_zip_dobject_payloads(source):
            if object_name == DIFOP_OBJECT:
                difop_records += 1
                if angles is None:
                    try:
                        angles = parse_difop_angles(parse_difop_payload(payload).raw)
                    except (EOFError, ValueError):
                        pass
                continue
            if object_name != MSOP_OBJECT:
                continue
            batch = parse_msop_batch_payload(payload)
            msop_batches += 1
            for item in batch.packets:
                station = lookup(item.host_receive_utc_ticks)
                if station is None:
                    continue
                packet_key = (batch.session_id, item.sequence)
                station_seen = seen_packets.setdefault(station, set())
                if packet_key in station_seen:
                    duplicates += 1
                    continue
                station_seen.add(packet_key)
                packets.setdefault(station, {})[packet_key] = (
                    item.host_receive_utc_ticks,
                    item.sequence,
                    item.raw,
                )
                selected_packets += 1
        if angles is None:
            raise RuntimeError(f"no valid H32 DIFOP angles found before decoding {source}")
        written = 0
        for station, items in packets.items():
            count = save_frames(
                station,
                items,
                export_root,
                angles,
                source,
                frame_stride=frame_stride,
                seen_frame_keys=seen_frames,
            )
            frame_counts[station] = frame_counts.get(station, 0) + count
            written += count
        source_stats.append(
            {
                "source": str(source.resolve()),
                "msop_batches": msop_batches,
                "difop_records": difop_records,
                "selected_packets": selected_packets,
                "duplicate_packets": duplicates,
                "frames_written": written,
            }
        )
        print(f"  selected_packets={selected_packets} frames={written} duplicates={duplicates}", flush=True)
    return {"frame_counts": frame_counts, "sources": source_stats}


def parse_rtk_sources(paths: list[Path], parsed_root: Path) -> dict[str, object]:
    rows = []
    source_stats = []
    for path in paths:
        capture_rows = parse_rtk_capture(
            read_capture(path),
            accepted_prefixes=("$GNGGA", "$GPGGA", "$GNHPR"),
        )
        for row in capture_rows:
            row["capture_source"] = str(path.resolve())
        rows.extend(capture_rows)
        source_stats.append(
            {
                "source": str(path.resolve()),
                "rows": len(capture_rows),
                "gga_valid": sum(row.get("type") == "GGA" and row.get("checksum_valid") for row in capture_rows),
                "gnhpr_valid": sum(
                    row.get("type") == "GNHPR" and row.get("checksum_valid") and row.get("heading_valid")
                    for row in capture_rows
                ),
            }
        )
    parsed_root.mkdir(parents=True, exist_ok=True)
    write_jsonl(parsed_root / "rtk.jsonl", rows)
    write_jsonl(parsed_root / "imu.jsonl", [])
    return {"rows": len(rows), "sources": source_stats}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--segments-csv", type=Path, required=True)
    parser.add_argument("--lidar-dlog", type=Path, action="append", default=[])
    parser.add_argument("--rtk-rscap", type=Path, action="append", required=True)
    parser.add_argument("--out", type=Path, required=True)
    parser.add_argument("--expected-stations", type=int, default=0)
    parser.add_argument("--frame-stride", type=int, default=5)
    parser.add_argument("--rtk-max-dt-ms", type=float, default=200.0)
    parser.add_argument("--reuse-export", action="store_true", help="Keep existing export/ and resume parsed/combined stages.")
    parser.add_argument("--overwrite", action="store_true")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    if args.frame_stride < 1:
        raise SystemExit("--frame-stride must be >= 1")
    if not args.reuse_export and not args.lidar_dlog:
        raise SystemExit("at least one --lidar-dlog is required unless --reuse-export is used")
    for source in [args.segments_csv, *args.lidar_dlog, *args.rtk_rscap]:
        if not source.is_file():
            raise FileNotFoundError(source)
    if args.reuse_export:
        export_root = args.out / "export"
        if not export_root.is_dir():
            raise FileNotFoundError(f"--reuse-export requested but missing {export_root}")
        for name in ("parsed", "combined", "export_summary.json"):
            target = args.out / name
            if target.is_dir():
                shutil.rmtree(target)
            elif target.exists():
                target.unlink()
    elif args.out.exists() and any(args.out.iterdir()):
        if not args.overwrite:
            raise FileExistsError(f"{args.out} is non-empty; pass --overwrite")
        for name in ("export", "parsed", "combined", "export_summary.json"):
            target = args.out / name
            if target.is_dir():
                shutil.rmtree(target)
            elif target.exists():
                target.unlink()
    args.out.mkdir(parents=True, exist_ok=True)
    windows = load_windows(args.segments_csv)
    expected_ids = sorted({item.station_id for item in windows})
    if args.reuse_export:
        frame_counts = {
            station.name: len(list((station / "frames").glob("*.npz")))
            for station in (args.out / "export").iterdir()
            if station.is_dir()
        }
        lidar_summary = {"frame_counts": frame_counts, "sources": [], "reused_export": True}
    else:
        lidar_summary = scan_dlog_sources(
            args.lidar_dlog,
            windows,
            args.out / "export",
            frame_stride=args.frame_stride,
        )
        frame_counts = lidar_summary["frame_counts"]
    exported_ids = sorted(station for station, count in frame_counts.items() if count)
    missing = sorted(set(expected_ids) - set(exported_ids))
    if missing:
        raise RuntimeError(f"stations without decoded H32 frames: {missing}")
    if args.expected_stations and len(exported_ids) != args.expected_stations:
        raise RuntimeError(f"expected {args.expected_stations} stations, exported {len(exported_ids)}")
    rtk_summary = parse_rtk_sources(args.rtk_rscap, args.out / "parsed")
    lidar_segments = [(station, args.out / "export" / station / "frames") for station in exported_ids]
    combined_summary = build_combined(
        lidar_segments,
        [args.out / "parsed" / "rtk.jsonl"],
        [],
        args.out / "combined",
        rtk_max_dt_ms=args.rtk_max_dt_ms,
        time_basis="host",
        overwrite=True,
    )
    summary = {
        "role": "G90 GNGGA/GNHPR + H32 DLog static-window export",
        "segments_csv": str(args.segments_csv.resolve()),
        "time_basis": "H32 MSOP host_receive_utc_ticks <-> G90 rscap host_receive_utc_ns",
        "imu_used": False,
        "expected_station_ids": expected_ids,
        "station_count": len(exported_ids),
        "lidar": lidar_summary,
        "rtk": rtk_summary,
        "combined": combined_summary,
        "outputs": {
            "combined": str((args.out / "combined").resolve()),
            "manifest": str((args.out / "combined" / "manifest.csv").resolve()),
        },
    }
    (args.out / "export_summary.json").write_text(
        json.dumps(summary, ensure_ascii=False, indent=2) + "\n",
        encoding="utf-8",
    )
    print(json.dumps({"stations": len(exported_ids), "combined": combined_summary}, ensure_ascii=False, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())