#!/usr/bin/env python3
"""
stream_gps.py - Mode 1: continuous real-time positioning for the u-blox ZED-X20P.

Best-available correction tier is selected automatically:
  1. Network RTK via NTRIP  (default: Mississippi GCGC RTN, free registration at
     http://rtn.usm.edu/RegisterAccount.aspx) -> onboard RTK engine, 1-3 cm.
     Credentials: env GCGC_USER / GCGC_PASS (or --ntrip-user/--ntrip-pass).
  2. Galileo HAS (E6-B, no internet needed)  -> ~20-25 cm.  Requires firmware
     HPG >= 2.10 (auto-probed via MON-VER; this project's module reports 2.10).
  3. Neither available -> the onboard multiband+SBAS solution (~1 m) cannot be
     improved in real time; per project policy the enhanced mode is SKIPPED
     (run with --log-only to still record a .ubx for post-processing).

Every session simultaneously logs raw RXM-RAWX/SFRBX to a .ubx file, so a
streaming session can always be post-processed later with process_gps.py.

Buoy deployment defaults: the receiver's dynamic platform model is set to "sea"
(--dynmodel to override) and the NTRIP GGA uplink sends the LIVE position so the
VRS reference follows platform drift (--gga-fixed to disable for a fixed antenna).

Outputs: live console status, stream_<ts>.csv (decimal-degree lat/lon),
         stream_<ts>.ubx (raw capture).

Usage:
    venv/bin/python stream_gps.py --port /dev/ttyUSB0            (Linux)
    venv/bin/python stream_gps.py --port COM7                    (Windows)
    venv/bin/python stream_gps.py --replay capture.ubx           (offline test)

Datum note: NTRIP corrections define the output frame. GCGC RTN broadcasts
NAD83(2011) epoch 2010.00 - constant ~1.5 m from WGS84/ITRF in CONUS; the CSV
header records which tier (and therefore frame) was active.
"""

import argparse
import csv
import os
import queue
import signal
import socket
import sys
import threading
import time
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path

from pyubx2 import UBXMessage, UBXReader, UBX_PROTOCOL, SET_LAYER_RAM

FIX_NAME = {0: "none", 1: "DR", 2: "2D", 3: "3D", 4: "GNSS+DR", 5: "time-only"}
CARR_NAME = {0: "", 1: "RTK-FLOAT", 2: "RTK-FIXED"}

# UBX-CFG keys (u-blox X20 HPG interface descriptions; NAVCOR from HPG 2.10 JSON)
K = {
    "UART1_BAUD":        0x40520001,
    "UART1_IN_RTCM3":    0x10730004,
    "USB_IN_RTCM3":      0x10770004,
    "PVT_UART1":         0x20910007, "PVT_USB":       0x20910009,
    "HPPOSLLH_UART1":    0x20910034, "HPPOSLLH_USB":  0x20910036,
    "RAWX_UART1":        0x209102a5, "RAWX_USB":      0x209102a7,
    "SFRBX_UART1":       0x20910232, "SFRBX_USB":     0x20910234,
    "NAVCOR_ENABLE_HOST":    0x100d0001,   # HPG >= 2.10
    "NAVCOR_ENABLE_GAL_HAS": 0x100d0002,   # HPG >= 2.10
    "DYNMODEL":          0x20110021,       # CFG-NAVSPG-DYNMODEL
}

# CFG-NAVSPG-DYNMODEL values; default "sea" for the buoy deployment
DYNMODEL = {"portable": 0, "stationary": 2, "pedestrian": 3, "automotive": 4,
            "sea": 5, "airborne1g": 6, "airborne2g": 7, "airborne4g": 8, "wrist": 9}


def log(msg):
    print(f"[stream_gps] {msg}", flush=True)


# ----------------------------------------------------------------------------
# Receiver configuration
# ----------------------------------------------------------------------------

def valset(ser, ubr, cfg, label):
    """Send CFG-VALSET (RAM) and wait briefly for ACK; returns True on ACK-ACK."""
    msg = UBXMessage.config_set(layers=SET_LAYER_RAM, transaction=0, cfgData=cfg)
    ser.write(msg.serialize())
    t_end = time.time() + 2.0
    while time.time() < t_end:
        try:
            _, parsed = ubr.read()
        except Exception:
            continue
        if parsed is None:
            continue
        if parsed.identity == "ACK-ACK":
            return True
        if parsed.identity == "ACK-NAK":
            log(f"  NAK for {label}")
            return False
    log(f"  no ACK for {label} (continuing)")
    return False


def poll_fw_version(ser, ubr):
    """Poll MON-VER; returns firmware string like 'HPG 2.10' or None."""
    ser.write(UBXMessage("MON", "MON-VER", 2).serialize())  # 2 = POLL msgmode
    t_end = time.time() + 3.0
    while time.time() < t_end:
        try:
            _, parsed = ubr.read()
        except Exception:
            continue
        if parsed is not None and parsed.identity == "MON-VER":
            sw = getattr(parsed, "swVersion", b"")
            sw = sw.decode(errors="replace") if isinstance(sw, bytes) else str(sw)
            return sw.strip("\x00 ")
    return None


def configure_receiver(ser, ubr, want_has, dynmodel):
    """Enable RTCM input, HP output, raw logging; optionally HAS. RAM layer only."""
    base_cfg = [
        (K["UART1_IN_RTCM3"], 1), (K["USB_IN_RTCM3"], 1),
        (K["PVT_UART1"], 1), (K["PVT_USB"], 1),
        (K["HPPOSLLH_UART1"], 1), (K["HPPOSLLH_USB"], 1),
        (K["RAWX_UART1"], 1), (K["RAWX_USB"], 1),
        (K["SFRBX_UART1"], 1), (K["SFRBX_USB"], 1),
        (K["DYNMODEL"], DYNMODEL[dynmodel]),
    ]
    ok = valset(ser, ubr, base_cfg, f"base output/input config (dynmodel={dynmodel})")
    if want_has:
        ok_has = valset(ser, ubr, [(K["NAVCOR_ENABLE_HOST"], 0),
                                   (K["NAVCOR_ENABLE_GAL_HAS"], 1)], "GAL HAS enable")
        return ok, ok_has
    return ok, False


# ----------------------------------------------------------------------------
# NTRIP
# ----------------------------------------------------------------------------

def caster_reachable(server, port, timeout=8):
    try:
        with socket.create_connection((server, port), timeout=timeout):
            return True
    except OSError:
        return False


class SerialRTCMSink:
    """Write-only wrapper handed to GNSSNTRIPClient as its output stream."""
    def __init__(self, ser):
        self._ser = ser
        self.bytes_out = 0

    def write(self, data):
        # GNSSNTRIPClient may hand us (raw, parsed) tuples or raw bytes
        if isinstance(data, tuple):
            data = data[0]
        if isinstance(data, (bytes, bytearray)):
            self._ser.write(bytes(data))
            self.bytes_out += len(data)

    # queue-style interface (pygnssutils uses .put on Queue outputs)
    def put(self, item):
        self.write(item)


def start_ntrip(args, sink, reflat, reflon, app=None):
    """app (with .get_coordinates()) => live GGA that follows a drifting platform;
    without app, falls back to a fixed reference position."""
    from pygnssutils import GNSSNTRIPClient
    client = GNSSNTRIPClient(app=app)
    kwargs = dict(
        server=args.ntrip_caster, port=args.ntrip_port, https=0,
        mountpoint=args.ntrip_mount, datatype="RTCM",
        ntripuser=args.ntrip_user, ntrippassword=args.ntrip_pass,
        ggainterval=args.gga_interval,
        ggamode=0 if app is not None else 1,    # 0 = live from receiver, 1 = fixed
        reflat=reflat, reflon=reflon, refalt=0.0, refsep=0.0,
        output=sink,
    )
    ok = client.run(**kwargs)
    return client if ok else None


# ----------------------------------------------------------------------------
# Main streaming loop
# ----------------------------------------------------------------------------

class Session:
    def __init__(self, args):
        self.args = args
        self.q = queue.Queue()
        self.stop = threading.Event()
        self.fix_hist = Counter()
        self.hacc = []
        self.tier = "none"
        # latest fix, served to GNSSNTRIPClient.get_coordinates() for live GGA
        self.live = {"lat": 0.0, "lon": 0.0, "alt": 0.0, "sep": 0.0}
        ts = datetime.now(timezone.utc).strftime("%Y%m%d_%H%M%S")
        self.csv_path = Path(f"stream_{ts}.csv")
        self.ubx_path = Path(f"stream_{ts}.ubx")

    def get_coordinates(self):
        """pygnssutils app interface: live position for the NTRIP GGA uplink,
        so the VRS reference follows a drifting platform (buoy)."""
        return dict(self.live)

    def reader_thread(self, stream, raw_log):
        """Serial/file -> raw .ubx tee -> parsed NAV messages -> queue."""
        ubr = UBXReader(stream, protfilter=UBX_PROTOCOL, quitonerror=0)
        while not self.stop.is_set():
            try:
                raw, parsed = ubr.read()
            except Exception:
                if self.args.replay:
                    break
                continue
            if raw is None:
                if self.args.replay:
                    break
                continue
            if raw_log:
                raw_log.write(raw)
            if parsed is None:
                continue
            if parsed.identity in ("NAV-PVT", "NAV-HPPOSLLH"):
                self.q.put(parsed)
        self.q.put(None)

    def run(self):
        args = self.args
        # ---------------- input: serial or replay file ----------------------
        if args.replay:
            stream = open(args.replay, "rb")
            ser = None
            log(f"REPLAY mode from {args.replay} (no config writes, no NTRIP)")
        else:
            import serial
            from serial.tools import list_ports
            if not args.port:
                ports = [p.device for p in list_ports.comports()]
                sys.exit(f"--port required. Detected ports: {ports or 'none'}")
            ser = serial.Serial(args.port, args.baud, timeout=1)
            stream = ser
            log(f"connected {args.port} @ {args.baud}")

            cfg_ubr = UBXReader(ser, protfilter=UBX_PROTOCOL, quitonerror=0)
            fw = poll_fw_version(ser, cfg_ubr)
            log(f"firmware: {fw or 'unknown'}")
            fw_has = bool(fw and any(
                fw.split("HPG")[-1].strip().startswith(v)
                for v in ("2.1", "2.2", "3.")))

            # ------------- correction tier selection -------------------------
            ntrip_ok = False
            if not args.log_only and args.ntrip_user and args.ntrip_pass:
                ntrip_ok = caster_reachable(args.ntrip_caster, args.ntrip_port)
                if not ntrip_ok:
                    log(f"NTRIP caster {args.ntrip_caster}:{args.ntrip_port} unreachable")
            elif not args.log_only:
                log("no NTRIP credentials (set GCGC_USER/GCGC_PASS or --ntrip-user/-pass)")

            want_has = not ntrip_ok and not args.log_only and fw_has
            if ntrip_ok:
                self.tier = f"NTRIP RTK ({args.ntrip_caster}/{args.ntrip_mount})"
            elif want_has:
                self.tier = "Galileo HAS (E6-B, ~20 cm)"
            elif args.log_only:
                self.tier = "log-only (onboard solution)"
            else:
                log("no correction source available that improves on the onboard "
                    "solution (NTRIP unreachable/no credentials; HAS needs FW >= "
                    "HPG 2.10). Skipping enhanced streaming per project policy.")
                log("use --log-only to record raw data for post-processing anyway.")
                ser.close()
                return 2

            ok_base, ok_has = configure_receiver(ser, cfg_ubr, want_has, args.dynmodel)
            log(f"receiver configured (base={'ok' if ok_base else 'partial'}"
                + (f", HAS={'ok' if ok_has else 'failed'}" if want_has else "") + ")")

        log(f"correction tier: {self.tier}")

        # ---------------- outputs -------------------------------------------
        raw_log = None if args.replay else open(self.ubx_path, "wb")
        csv_f = open(self.csv_path, "w", newline="")
        csv_f.write(f"# stream_gps.py session - correction tier: {self.tier}\n")
        csv_f.write("# datum: NAD83(2011) if GCGC NTRIP tier, else WGS84/ITRF\n")
        writer = csv.writer(csv_f)
        writer.writerow(["utc", "lat_deg", "lon_deg", "ellip_height_m",
                         "hAcc_m", "vAcc_m", "fixType", "carrSoln", "numSV"])

        # ---------------- threads -------------------------------------------
        t_read = threading.Thread(target=self.reader_thread, args=(stream, raw_log),
                                  daemon=True)
        t_read.start()

        ntrip_client, sink = None, None
        if ser is not None and self.tier.startswith("NTRIP"):
            # need a first fix for the GGA reference position
            reflat, reflon = args.lat, args.lon
            if reflat is None:
                log("waiting for first fix to seed NTRIP GGA position ...")
                first = self._wait_first_fix(timeout=60)
                if first is None:
                    log("no fix within 60 s; using 0/0 GGA (VRS may reject)")
                    reflat = reflon = 0.0
                else:
                    reflat, reflon = first
            sink = SerialRTCMSink(ser)
            app = None if args.gga_fixed else self
            log("GGA uplink: " + ("fixed reference" if app is None
                                  else "live (follows platform drift)"))
            ntrip_client = start_ntrip(args, sink, reflat, reflon, app=app)
            if ntrip_client is None:
                log("NTRIP connection failed (check credentials/mountpoint) - "
                    "continuing with onboard solution only")
                self.tier += " [FAILED - onboard only]"

        signal.signal(signal.SIGINT, lambda *_: self.stop.set())

        # ---------------- consume + display ---------------------------------
        last_pvt = {}
        try:
            while not self.stop.is_set():
                try:
                    m = self.q.get(timeout=1.0)
                except queue.Empty:
                    continue
                if m is None:
                    break
                if m.identity == "NAV-PVT":
                    last_pvt = dict(
                        utc=datetime(m.year, m.month, m.day, m.hour, m.min, m.second,
                                     tzinfo=timezone.utc),
                        fixType=m.fixType, carrSoln=m.carrSoln, numSV=m.numSV,
                        lat=m.lat, lon=m.lon, height=m.height / 1000.0,
                        hAcc=m.hAcc / 1000.0, vAcc=m.vAcc / 1000.0)
                elif m.identity == "NAV-HPPOSLLH" and last_pvt:
                    # high-precision refinement: lat/lon + Hp components (deg)
                    last_pvt["lat"] = m.lat + m.latHp
                    last_pvt["lon"] = m.lon + m.lonHp
                    last_pvt["height"] = (m.height + m.heightHp) / 1000.0
                    last_pvt["hAcc"] = m.hAcc / 1000.0
                    last_pvt["vAcc"] = m.vAcc / 1000.0
                    continue
                if not last_pvt:
                    continue
                p = last_pvt
                self.live.update(lat=p["lat"], lon=p["lon"], alt=p["height"])
                self.fix_hist[(p["fixType"], p["carrSoln"])] += 1
                self.hacc.append(p["hAcc"])
                writer.writerow([p["utc"].isoformat(), f"{p['lat']:.9f}",
                                 f"{p['lon']:.9f}", f"{p['height']:.3f}",
                                 f"{p['hAcc']:.3f}", f"{p['vAcc']:.3f}",
                                 p["fixType"], p["carrSoln"], p["numSV"]])
                n_epochs = sum(self.fix_hist.values())
                if not sys.stdout.isatty() and n_epochs % 60 != 1:
                    continue                      # avoid log spam when piped
                fixlab = FIX_NAME.get(p["fixType"], "?")
                if p["carrSoln"]:
                    fixlab += "/" + CARR_NAME[p["carrSoln"]]
                rtcm = f"  rtcm {sink.bytes_out//1024} kB" if sink else ""
                end = "   " if sys.stdout.isatty() else "\n"
                status = (f"\r{p['utc']:%H:%M:%S}  {fixlab:<12s} "
                          f"{p['lat']:+.9f} {p['lon']:+.9f}  h={p['height']:8.3f} m  "
                          f"hAcc={p['hAcc']:6.3f} m  sv={p['numSV']:2d}{rtcm}{end}")
                sys.stdout.write(status)
                sys.stdout.flush()
        finally:
            self.stop.set()
            print()
            if ntrip_client is not None:
                try:
                    ntrip_client.stop()
                except Exception:
                    pass
            csv_f.close()
            if raw_log:
                raw_log.close()
            if ser:
                ser.close()
            self._summary()
        return 0

    def _wait_first_fix(self, timeout):
        t_end = time.time() + timeout
        while time.time() < t_end:
            try:
                m = self.q.get(timeout=1.0)
            except queue.Empty:
                continue
            if m is not None and m.identity == "NAV-PVT" and m.fixType >= 2:
                self.live.update(lat=m.lat, lon=m.lon, alt=m.height / 1000.0)
                return (m.lat, m.lon)
        return None

    def _summary(self):
        n = sum(self.fix_hist.values())
        if not n:
            log("no position epochs received")
            return
        log(f"session summary: {n} epochs, tier = {self.tier}")
        for (fix, carr), c in sorted(self.fix_hist.items(), key=lambda kv: -kv[1]):
            lab = FIX_NAME.get(fix, "?") + ("/" + CARR_NAME[carr] if carr else "")
            log(f"  {lab:<14s} {c:6d}  ({100*c/n:.1f}%)")
        h = sorted(self.hacc)
        log(f"  hAcc: mean {sum(h)/len(h):.3f} m, median {h[len(h)//2]:.3f} m, "
            f"95% {h[int(len(h)*0.95)]:.3f} m")
        if not self.args.replay:
            log(f"  outputs: {self.csv_path}, {self.ubx_path}")
            log(f"  post-process the raw log: venv/bin/python process_gps.py {self.ubx_path}")


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                 formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--port", help="serial port (COM7, /dev/ttyUSB0, ...)")
    ap.add_argument("--baud", type=int, default=38400)
    ap.add_argument("--replay", type=Path, help="read a .ubx file instead of serial "
                                                "(offline parser/CSV test mode)")
    ap.add_argument("--log-only", action="store_true",
                    help="skip corrections; just log raw data + onboard positions")
    ap.add_argument("--lat", type=float, help="initial latitude for NTRIP GGA (else first fix)")
    ap.add_argument("--lon", type=float, help="initial longitude for NTRIP GGA")
    ap.add_argument("--gga-fixed", action="store_true",
                    help="send a FIXED GGA reference instead of live position "
                         "(only for a truly stationary antenna; default is live, "
                         "so the VRS reference follows a drifting buoy)")
    ap.add_argument("--dynmodel", default="sea", choices=sorted(DYNMODEL),
                    help="receiver dynamic platform model, set in RAM at startup "
                         "(default: sea, for buoy deployment)")
    ap.add_argument("--ntrip-caster", default="rtn.usm.edu")
    ap.add_argument("--ntrip-port", type=int, default=2101)
    ap.add_argument("--ntrip-mount", default="VRS_GNSS_RTCM34_NAD83")
    ap.add_argument("--ntrip-user", default=os.environ.get("GCGC_USER"))
    ap.add_argument("--ntrip-pass", default=os.environ.get("GCGC_PASS"))
    ap.add_argument("--gga-interval", type=int, default=15)
    args = ap.parse_args()

    if not args.replay and not args.port:
        ap.error("one of --port or --replay is required")
    sys.exit(Session(args).run())


if __name__ == "__main__":
    main()