#!/usr/bin/env python3
from __future__ import annotations

import os
import time
import json
import socket
import serial
import argparse
import signal
import struct
from pathlib import Path
from typing import Optional, Dict, Tuple, List


# ============================================================
# Defaults
# ============================================================

ZED_PORT_DEFAULT = "/dev/serial/by-id/usb-u-blox_AG_-_www.u-blox.com_u-blox_GNSS_receiver-if00"
ZED_BAUD_DEFAULT = 115200

LORA_PORT_DEFAULT = "/dev/serial/by-id/usb-1a86_USB_Serial-if00-port0"
LORA_BAUD_DEFAULT = 9600  # <-- wichtig: LoRa UART ist bei dir 9600

STATE_PATH_DEFAULT = "/run/rtk/state.json"

SVIN_TARGET_DUR_DEFAULT = 120
SVIN_TARGET_ACC_DEFAULT = 11.0

# Read/Loop tuning
READ_IDLE_SLEEP_S = 0.002
STATE_WRITE_INTERVAL_S = 1.0

HOSTNAME = socket.gethostname()


# ============================================================
# State / Files
# ============================================================

def ensure_dir(path: str):
    Path(path).mkdir(parents=True, exist_ok=True)


def ensure_parent_dir(file_path: str):
    parent = os.path.dirname(file_path) or "."
    ensure_dir(parent)


def write_state(path: str, data: dict):
    ensure_parent_dir(path)
    tmp = path + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(data, f, ensure_ascii=False)
    os.replace(tmp, path)


def now_ts() -> str:
    return time.strftime("%Y-%m-%d %H:%M:%S")


# ============================================================
# Serial
# ============================================================

def open_serial(port: str, baud: int) -> serial.Serial:
    print(f"🔌 Öffne {port} @ {baud} …")
    s = serial.Serial(port, baud, timeout=0)  # non-blocking
    try:
        s.dtr = True
        s.rts = False
    except Exception:
        pass
    try:
        s.reset_input_buffer()
    except Exception:
        pass
    print("✅ Port offen.")
    return s


# ============================================================
# UBX helpers
# ============================================================

def ubx_checksum(payload: bytes) -> bytes:
    ck_a = 0
    ck_b = 0
    for b in payload:
        ck_a = (ck_a + b) & 0xFF
        ck_b = (ck_b + ck_a) & 0xFF
    return bytes([ck_a, ck_b])


def ubx_build(msg_class: int, msg_id: int, payload: bytes) -> bytes:
    length = len(payload)
    hdr = bytes([0xB5, 0x62, msg_class & 0xFF, msg_id & 0xFF, length & 0xFF, (length >> 8) & 0xFF])
    body = bytes([msg_class & 0xFF, msg_id & 0xFF, length & 0xFF, (length >> 8) & 0xFF]) + payload
    cks = ubx_checksum(body)
    return hdr + payload + cks


def ubx_send(ser: serial.Serial, msg_class: int, msg_id: int, payload: bytes = b""):
    ser.write(ubx_build(msg_class, msg_id, payload))


def ubx_cfg_msg(ser: serial.Serial, msg_class: int, msg_id: int, rate_usb: int):
    """
    UBX-CFG-MSG (0x06 0x01) legacy: payload = class, id, rateI2C, rateUART1, rateUART2, rateUSB, rateSPI
    """
    payload = bytes([msg_class & 0xFF, msg_id & 0xFF, 0x00, 0x00, 0x00, rate_usb & 0xFF, 0x00])
    ubx_send(ser, 0x06, 0x01, payload)


# ============================================================
# RTCM Tap (binary file + rotation)
# ============================================================

class RtcmTap:
    def __init__(self, path: Optional[str], rotate_mb: int = 0, max_files: int = 5):
        self.path = path
        self.rotate_bytes = int(rotate_mb) * 1024 * 1024 if rotate_mb and rotate_mb > 0 else 0
        self.max_files = max(1, int(max_files))
        self._fh = None
        self.ok = False
        self.bytes_written = 0

        if self.path:
            self._open_append()

    def _open_append(self):
        try:
            ensure_parent_dir(self.path)  # type: ignore[arg-type]
            self._fh = open(self.path, "ab", buffering=0)
            self.ok = True
        except Exception as e:
            print(f"⚠️  RTCM tap open failed ({self.path}): {e}")
            self._fh = None
            self.ok = False

    def _rotate_if_needed(self):
        if not self.path or not self.rotate_bytes or not self.ok:
            return
        try:
            size = os.path.getsize(self.path)
        except Exception:
            return
        if size < self.rotate_bytes:
            return

        try:
            if self._fh:
                self._fh.close()
        except Exception:
            pass
        self._fh = None

        for i in range(self.max_files - 1, 0, -1):
            src = f"{self.path}.{i}"
            dst = f"{self.path}.{i+1}"
            if os.path.exists(src):
                try:
                    if os.path.exists(dst):
                        os.remove(dst)
                except Exception:
                    pass
                try:
                    os.rename(src, dst)
                except Exception:
                    pass

        try:
            dst1 = f"{self.path}.1"
            if os.path.exists(dst1):
                os.remove(dst1)
            os.rename(self.path, dst1)
            print(f"🌀 RTCM tap rotated: {self.path} -> {dst1}")
        except Exception as e:
            print(f"⚠️  RTCM tap rotate failed: {e}")

        self._open_append()

    def write(self, frame: bytes):
        if not self.path or not self.ok or self._fh is None:
            return
        try:
            self._fh.write(frame)
            self.bytes_written += len(frame)
        except Exception as e:
            print(f"⚠️  RTCM tap write failed: {e}")
            self.ok = False
            try:
                self._fh.close()
            except Exception:
                pass
            self._fh = None
            return

        self._rotate_if_needed()

    def close(self):
        try:
            if self._fh:
                self._fh.close()
        except Exception:
            pass
        self._fh = None


# ============================================================
# RTCM3 CRC24Q
# ============================================================

_CRC24Q_POLY = 0x1864CFB

def crc24q(data: bytes) -> int:
    crc = 0
    for b in data:
        crc ^= (b << 16)
        for _ in range(8):
            crc <<= 1
            if crc & 0x1000000:
                crc ^= _CRC24Q_POLY
    return crc & 0xFFFFFF


def rtcm_msg_type(payload: bytes) -> Optional[int]:
    if len(payload) < 2:
        return None
    return (payload[0] << 4) | (payload[1] >> 4)


# ============================================================
# Stream Demux: RTCM + UBX from the same serial stream
# ============================================================

class StreamDemux:
    """
    Extract RTCM3 frames (0xD3...) AND UBX frames (0xB5 0x62...) from a mixed stream.
    Important: we must NOT discard UBX while searching for RTCM.
    """

    def __init__(self):
        self.buf = bytearray()
        self.rtcm_invalid_crc = 0
        self.ubx_invalid_crc = 0

    def feed(self, data: bytes):
        if data:
            self.buf.extend(data)

    @staticmethod
    def _try_parse_rtcm(buf: bytearray, i: int) -> Optional[Tuple[int, bytes]]:
        if len(buf) - i < 6:
            return None
        if buf[i] != 0xD3:
            return None
        length = ((buf[i+1] & 0x03) << 8) | buf[i+2]
        if length > 1023:
            return None
        need = 3 + length + 3
        if len(buf) - i < need:
            return None
        frame = bytes(buf[i:i+need])
        body = frame[:3+length]
        crc_bytes = frame[3+length:3+length+3]
        want = (crc_bytes[0] << 16) | (crc_bytes[1] << 8) | crc_bytes[2]
        got = crc24q(body)
        if got != want:
            return (-need, b"")  # signal: parsed length, but bad crc
        return (need, frame)

    @staticmethod
    def _try_parse_ubx(buf: bytearray, i: int) -> Optional[Tuple[int, int, int, bytes]]:
        if len(buf) - i < 8:
            return None
        if buf[i] != 0xB5 or buf[i+1] != 0x62:
            return None
        msg_class = buf[i+2]
        msg_id = buf[i+3]
        ln = buf[i+4] | (buf[i+5] << 8)
        need = 6 + ln + 2
        if len(buf) - i < need:
            return None
        payload = bytes(buf[i+6:i+6+ln])
        ck_rx = bytes(buf[i+6+ln:i+6+ln+2])
        ck_calc = ubx_checksum(bytes([msg_class, msg_id, ln & 0xFF, (ln >> 8) & 0xFF]) + payload)
        if ck_calc != ck_rx:
            return (-need, msg_class, msg_id, b"")  # bad crc
        return (need, msg_class, msg_id, payload)

    def pop_frames(self) -> Tuple[List[bytes], List[Tuple[int, int, bytes]]]:
        """
        Returns (rtcm_frames_valid, ubx_frames_valid)
        ubx_frames_valid items are (cls, id, payload)
        """
        rtcm_out: List[bytes] = []
        ubx_out: List[Tuple[int, int, bytes]] = []

        i = 0
        while i < len(self.buf):
            b = self.buf[i]

            if b == 0xD3:
                r = self._try_parse_rtcm(self.buf, i)
                if r is None:
                    break
                n, frame = r
                if n < 0:
                    self.rtcm_invalid_crc += 1
                    i += (-n)
                else:
                    rtcm_out.append(frame)
                    i += n
                continue

            if b == 0xB5:
                r2 = self._try_parse_ubx(self.buf, i)
                if r2 is None:
                    break
                n, cls, mid, payload = r2
                if n < 0:
                    self.ubx_invalid_crc += 1
                    i += (-n)
                else:
                    ubx_out.append((cls, mid, payload))
                    i += n
                continue

            i += 1

        if i > 0:
            del self.buf[:i]
        if len(self.buf) > 1024 * 1024:
            self.buf = self.buf[-1024*1024:]

        return rtcm_out, ubx_out


# ============================================================
# NAV-SVIN parsing (ZED-F9P)
# ============================================================

import struct

def parse_nav_svin(payload: bytes):
    # NAV-SVIN (0x01 0x3B) has a 4-byte header:
    #   uint8  version
    #   uint8[3] reserved0
    # then the actual fields start at offset 4.   [oai_citation:1‡docs.ros.org](https://docs.ros.org/en/kinetic/api/ublox_msgs/html/msg/NavSVIN.html)
    if len(payload) < 40:
        return None

    version = payload[0]
    # payload[1:4] reserved0

    off = 4
    iTOW, dur, meanX, meanY, meanZ = struct.unpack_from("<IIiii", payload, off)
    off += 20

    meanXHP, meanYHP, meanZHP, _reserved1 = struct.unpack_from("<bbbB", payload, off)
    off += 4

    meanAcc, obs = struct.unpack_from("<II", payload, off)
    off += 8

    valid, active = struct.unpack_from("<BB", payload, off)
    # payload[off+2:off+4] reserved3

    # meanX/Y/Z in cm, HP in 0.1 mm -> meters
    meanX_m = (meanX * 0.01) + (meanXHP * 0.0001)
    meanY_m = (meanY * 0.01) + (meanYHP * 0.0001)
    meanZ_m = (meanZ * 0.01) + (meanZHP * 0.0001)

    meanAcc_m = meanAcc * 0.0001  # 0.1 mm -> m   [oai_citation:2‡docs.ros.org](https://docs.ros.org/en/kinetic/api/ublox_msgs/html/msg/NavSVIN.html)

    return {
        "version": int(version),
        "iTOW_ms": int(iTOW),
        "dur_s": int(dur),
        "meanX_m": float(meanX_m),
        "meanY_m": float(meanY_m),
        "meanZ_m": float(meanZ_m),
        "meanAcc_m": float(meanAcc_m),
        "obs": int(obs),
        "valid": bool(valid),
        "active": bool(active),
    }

# ============================================================
# Args / signals
# ============================================================

def parse_args():
    ap = argparse.ArgumentParser(
        description="ZED-F9P RTK Base: RTCM forward (USB -> LoRa) + state.json + RTCM tap"
    )
    ap.add_argument("--zed-port", default=ZED_PORT_DEFAULT)
    ap.add_argument("--zed-baud", type=int, default=ZED_BAUD_DEFAULT)
    ap.add_argument("--lora-port", default=LORA_PORT_DEFAULT)
    ap.add_argument("--lora-baud", type=int, default=LORA_BAUD_DEFAULT)
    ap.add_argument("--state-path", default=STATE_PATH_DEFAULT)
    ap.add_argument("--svin-target-dur", type=int, default=SVIN_TARGET_DUR_DEFAULT)
    ap.add_argument("--svin-target-acc", type=float, default=SVIN_TARGET_ACC_DEFAULT)

    # RTCM tap
    ap.add_argument(
        "--rtcm-tap-file",
        default="",
        help="Write CRC-valid RTCM frames to this binary file (append). Example: /run/rtk/rtcm.tap",
    )
    ap.add_argument("--rtcm-tap-rotate-mb", type=int, default=0)
    ap.add_argument("--rtcm-tap-max-files", type=int, default=5)

    # SVIN handling
    ap.add_argument(
        "--enable-svin-status",
        action="store_true",
        help="Enable NAV-SVIN output on USB and parse it into state.json.",
    )

    return ap.parse_args()


_STOP = False

def _handle_stop(signum, frame):
    global _STOP
    _STOP = True


# ============================================================
# Main
# ============================================================

def main():
    global _STOP
    args = parse_args()

    signal.signal(signal.SIGINT, _handle_stop)
    signal.signal(signal.SIGTERM, _handle_stop)

    start_time = time.time()

    zed = open_serial(args.zed_port, args.zed_baud)

    lora: Optional[serial.Serial] = None
    lora_ok = False
    try:
        lora = open_serial(args.lora_port, args.lora_baud)
        lora_ok = True
    except Exception as e:
        print(f"⚠️  LoRa nicht verfügbar ({e}). Forwarding deaktiviert.")

    # Enable NAV-SVIN output on USB (best-effort)
    if args.enable_svin_status:
        try:
            print("🛰️  Aktiviere UBX-NAV-SVIN auf USB (CFG-MSG)…")
            ubx_cfg_msg(zed, 0x01, 0x3B, 1)  # NAV-SVIN on USB
        except Exception as e:
            print(f"⚠️  NAV-SVIN enable failed: {e}")

    tap_path = args.rtcm_tap_file.strip() or None
    tap = RtcmTap(tap_path, rotate_mb=args.rtcm_tap_rotate_mb, max_files=args.rtcm_tap_max_files)
    if tap_path:
        print(f"🧷 RTCM tap aktiv: {tap_path} (rotate={args.rtcm_tap_rotate_mb}MB, keep={args.rtcm_tap_max_files})")

    # Runtime stats (VALID RTCM only)
    rtcm_frames = 0
    rtcm_bytes = 0
    lora_tx_bytes = 0
    last_rtcm_ts: Optional[float] = None
    last_state_write = 0.0

    # Diagnostics
    msg_counts: Dict[str, int] = {}

    # SVIN state
    svin_state = "UNKNOWN"
    svin_active = False
    svin_valid = False
    svin_dur_s = 0
    svin_mean_acc_m = 0.0

    demux = StreamDemux()

    print("📡 Starte RTCM Forwarding (ZED -> LoRa) …")
    print("ℹ️  Hinweis: Es werden NUR CRC-validierte RTCM3 Frames gezählt/weitergeleitet (und ggf. getappt).")

    while not _STOP:
        n = zed.in_waiting
        if n:
            chunk = zed.read(n)
            demux.feed(chunk)

            rtcm_list, ubx_list = demux.pop_frames()

            # Handle UBX frames (NAV-SVIN)
            if args.enable_svin_status and ubx_list:
                for cls, mid, payload in ubx_list:
                    if (cls, mid) == (0x01, 0x3B):  # NAV-SVIN
                        s = parse_nav_svin(payload)
                        if s:
                            svin_active = bool(int(s["active"]) != 0)
                            svin_valid = bool(int(s["valid"]) != 0)
                            svin_dur_s = int(s["dur_s"])
                            svin_mean_acc_m = float(s["meanAcc_m"])
                            if svin_active and not svin_valid:
                                svin_state = "SURVEY_IN"
                            elif svin_valid:
                                svin_state = "VALID"
                            else:
                                svin_state = "IDLE"

            # Handle RTCM frames
            for frame in rtcm_list:
                rtcm_frames += 1
                rtcm_bytes += len(frame)
                last_rtcm_ts = time.time()

                if tap_path and tap.ok:
                    tap.write(frame)

                length = ((frame[1] & 0x03) << 8) | frame[2]
                payload = frame[3: 3 + length]
                mt = rtcm_msg_type(payload)
                if mt is not None:
                    k = str(mt)
                    msg_counts[k] = msg_counts.get(k, 0) + 1

                if lora_ok and lora is not None:
                    try:
                        lora.write(frame)
                        lora_tx_bytes += len(frame)
                    except Exception as e:
                        print(f"⚠️  LoRa write failed -> deaktiviert ({e})")
                        lora_ok = False

        else:
            time.sleep(READ_IDLE_SLEEP_S)

        now = time.time()
        if now - last_state_write >= STATE_WRITE_INTERVAL_S:
            last_state_write = now
            rtcm_last_s = 9999.0 if last_rtcm_ts is None else round(now - last_rtcm_ts, 2)

            state = {
                "role": "base",
                "host": HOSTNAME,
                "zed_port": args.zed_port,
                "lora_port": args.lora_port,
                "uptime_s": int(now - start_time),

                "zed_ok": True,
                "lora_ok": bool(lora_ok),

                "svin_state": svin_state,
                "svin_active": bool(svin_active),
                "svin_valid": bool(svin_valid),
                "svin_dur_s": int(svin_dur_s),
                "svin_meanAcc_m": float(svin_mean_acc_m),
                "svin_target_dur_s": int(args.svin_target_dur),
                "svin_target_acc_m": float(args.svin_target_acc),

                "rtcm_out_frames": int(rtcm_frames),
                "rtcm_out_bytes": int(rtcm_bytes),
                "rtcm_last_s": float(rtcm_last_s),
                "lora_tx_bytes": int(lora_tx_bytes),

                "rtcm_invalid_crc": int(demux.rtcm_invalid_crc),
                "ubx_invalid_crc": int(demux.ubx_invalid_crc),
                "rtcm_msg_counts": msg_counts,

                "rtcm_tap_file": tap_path or "",
                "rtcm_tap_ok": bool(tap.ok) if tap_path else False,
                "rtcm_tap_bytes": int(tap.bytes_written) if tap_path else 0,

                "ts": now_ts(),
            }
            try:
                write_state(args.state_path, state)
            except Exception as e:
                print(f"⚠️  write_state failed: {e}")

    print("🛑 Stop signal received, exiting…")
    try:
        tap.close()
    except Exception:
        pass
    try:
        zed.close()
    except Exception:
        pass
    try:
        if lora is not None:
            lora.close()
    except Exception:
        pass


if __name__ == "__main__":
    main()