#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import argparse
import json
import os
import struct
import sys
import time
from dataclasses import dataclass, asdict
from typing import Optional, Tuple, List

import serial

# Optional: paho-mqtt (sudo apt install python3-paho-mqtt)
try:
    import paho.mqtt.client as mqtt
except Exception:
    mqtt = None


# =========================
# 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 send_ubx(ser: serial.Serial, cls: int, mid: int, payload: bytes = b"") -> None:
    head = b"\xB5\x62" + bytes([cls, mid]) + struct.pack("<H", len(payload))
    msg = bytes([cls, mid]) + struct.pack("<H", len(payload)) + payload
    ser.write(head + payload + ubx_checksum(msg))

def read_ubx_frame(ser: serial.Serial, timeout: float = 1.0) -> Optional[Tuple[int, int, bytes]]:
    end = time.time() + timeout
    # sync
    while time.time() < end:
        b = ser.read(1)
        if not b:
            continue
        if b == b"\xB5":
            b2 = ser.read(1)
            if b2 == b"\x62":
                break
    else:
        return None
    hdr = ser.read(4)
    if len(hdr) < 4:
        return None
    cls, mid, length = hdr[0], hdr[1], struct.unpack("<H", hdr[2:4])[0]
    payload = ser.read(length)
    cks = ser.read(2)
    if len(payload) != length or len(cks) != 2:
        return None
    if ubx_checksum(bytes([cls, mid]) + struct.pack("<H", length) + payload) != cks:
        return None
    return (cls, mid, payload)

def wait_ack(ser: serial.Serial, exp_cls: int, exp_id: int, timeout: float = 1.0) -> bool:
    end = time.time() + timeout
    while time.time() < end:
        f = read_ubx_frame(ser, timeout=end - time.time())
        if not f:
            continue
        c, m, p = f
        # ACK-ACK=0x05 0x01, ACK-NAK=0x05 0x00; payload: cls,id
        if c == 0x05 and m in (0x01, 0x00) and len(p) == 2 and (p[0], p[1]) == (exp_cls, exp_id):
            return (m == 0x01)
    return False

def mon_ver(ser: serial.Serial) -> Tuple[str, str]:
    send_ubx(ser, 0x0A, 0x04)  # MON-VER
    v1 = v2 = "?"
    end = time.time() + 0.7
    while time.time() < end:
        f = read_ubx_frame(ser, timeout=0.2)
        if not f:
            continue
        c, m, p = f
        if (c, m) != (0x0A, 0x04):
            continue
        for i in range(0, len(p), 30):
            s = p[i:i+30].split(b"\x00", 1)[0].decode("ascii", errors="ignore")
            if s.startswith(("ROM", "EXT", "PROT")):
                v1 = s
            if s.startswith("FWVER"):
                v2 = s
    return v1, v2


# =========================
# ZED-F9P config (minimal)
# =========================

def cfg_msg(ser: serial.Serial, msg_cls: int, msg_id: int, usb_rate: int) -> bool:
    # CFG-MSG 0x06 0x01: class,id, rates[UART1,UART2,USB,SPIF0,I2C,UART3]
    payload = bytes([msg_cls, msg_id, 0, 0, usb_rate, 0, 0])
    send_ubx(ser, 0x06, 0x01, payload)
    return wait_ack(ser, 0x06, 0x01, 1.0)

def quiet_nmea_on_usb(ser: serial.Serial) -> None:
    # NMEA class 0xF0 ids
    for mid in [0x00,0x01,0x02,0x03,0x04,0x05,0x06,0x07,0x08,0x09]:
        cfg_msg(ser, 0xF0, mid, 0)

def enable_nav_pvt_push_usb(ser: serial.Serial, rate: int = 1) -> bool:
    return cfg_msg(ser, 0x01, 0x07, rate)  # NAV-PVT

def enable_nav_svin_push_usb(ser: serial.Serial, rate: int = 1) -> bool:
    return cfg_msg(ser, 0x01, 0x3B, rate)  # NAV-SVIN

def poll_nav_svin(ser: serial.Serial, timeout: float = 1.0) -> Optional[Tuple[bool, bool, int, float]]:
    # poll NAV-SVIN
    send_ubx(ser, 0x01, 0x3B)
    frm = read_ubx_frame(ser, timeout=timeout)
    if not frm or (frm[0], frm[1]) != (0x01, 0x3B) or len(frm[2]) < 40:
        return None
    p = frm[2]

    # Layout-Variante A (bewährt bei dir): dur@8, meanAcc@28, valid@36, active@37
    dur = struct.unpack_from("<I", p, 8)[0]
    meanAcc_01mm = struct.unpack_from("<I", p, 28)[0]
    valid = p[36] != 0
    active = p[37] != 0
    return active, valid, dur, meanAcc_01mm / 10000.0

def poll_nav_pvt(ser: serial.Serial, timeout: float = 1.0) -> Optional[dict]:
    # poll NAV-PVT
    send_ubx(ser, 0x01, 0x07)
    frm = read_ubx_frame(ser, timeout=timeout)
    if not frm or (frm[0], frm[1]) != (0x01, 0x07) or len(frm[2]) < 92:
        return None
    p = frm[2]
    fixType = p[20]
    flags = p[21]
    carrSoln = (p[21] >> 6) & 0x03  # 0 none, 1 float, 2 fix
    numSV = p[23]
    lon = struct.unpack_from("<i", p, 24)[0] / 1e7
    lat = struct.unpack_from("<i", p, 28)[0] / 1e7
    height = struct.unpack_from("<i", p, 32)[0] / 1000.0
    hMSL = struct.unpack_from("<i", p, 36)[0] / 1000.0
    hAcc = struct.unpack_from("<I", p, 40)[0] / 1000.0
    vAcc = struct.unpack_from("<I", p, 44)[0] / 1000.0
    return {
        "fixType": int(fixType),
        "flags": int(flags),
        "carrSoln": int(carrSoln),
        "numSV": int(numSV),
        "lat": float(lat),
        "lon": float(lon),
        "height_m": float(height),
        "hMSL_m": float(hMSL),
        "hAcc_m": float(hAcc),
        "vAcc_m": float(vAcc),
    }

def fix_to_text(fixType: int, carrSoln: int) -> str:
    # fixType: 0 no, 1 dead reck, 2 2D, 3 3D, 4 GNSS+DR, 5 time only
    base = {0:"NO FIX",1:"DR",2:"2D",3:"3D",4:"GNSS+DR",5:"TIME"}.get(fixType, f"FIX{fixType}")
    if carrSoln == 1:
        return f"{base} / RTK-FLOAT"
    if carrSoln == 2:
        return f"{base} / RTK-FIX"
    return base


# =========================
# RTCM sniff + (optional) LoRa TX
# =========================

def rtcm_try_parse_stream(data: bytes, max_frames: int = 50):
    """
    Minimaler RTCM3 Frame Parser:
    - Preamble 0xD3
    - 10-bit length in next 2 bytes
    - frame = 3 header + length + 3 CRC
    Liefert: frames(list[bytes]), rest(bytes), bad_count(int), types(list[int])
    """
    i = 0
    frames = []
    bad = 0
    types = []
    n = len(data)

    def get_len(b1, b2):
        return ((b1 & 0x03) << 8) | b2

    while i + 3 <= n and len(frames) < max_frames:
        if data[i] != 0xD3:
            i += 1
            continue
        if i + 3 > n:
            break
        length = get_len(data[i+1], data[i+2])
        frame_len = 3 + length + 3
        if length > 2047 or frame_len < 6:
            bad += 1
            i += 1
            continue
        if i + frame_len > n:
            break  # rest kommt später
        frame = data[i:i+frame_len]
        # Type: first 12 bits after header => frame[3:5]
        if length >= 2:
            msg_type = ((frame[3] << 4) | (frame[4] >> 4)) & 0x0FFF
            types.append(int(msg_type))
        frames.append(frame)
        i += frame_len

    rest = data[i:] if i < n else b""
    return frames, rest, bad, types

def lora_send_frames(lora: serial.Serial, frames: List[bytes], max_payload: int):
    """
    Generisch: split frames in chunks <= max_payload and write raw.
    DX-LR03: je nach Firmware sind "transparent mode" oder AT+SEND nötig.
    -> Hier nur raw write. Falls du AT+SEND brauchst, sag Bescheid, dann baue ich das exakt für DX-LR03 ein.
    """
    sent = 0
    for fr in frames:
        off = 0
        while off < len(fr):
            chunk = fr[off:off+max_payload]
            lora.write(chunk)
            sent += len(chunk)
            off += len(chunk)
    return sent


# =========================
# MQTT
# =========================

class MqttPublisher:
    def __init__(self, host: str, port: int, base_topic: str, client_id: str,
                 username: Optional[str], password: Optional[str], retain: bool = True):
        if mqtt is None:
            raise RuntimeError("paho-mqtt fehlt. Install: sudo apt install python3-paho-mqtt")
        self.base_topic = base_topic.rstrip("/")
        self.retain = retain
        self.client = mqtt.Client(client_id=client_id, clean_session=True)
        if username:
            self.client.username_pw_set(username, password=password or None)

        # LWT: availability offline, wenn Prozess stirbt
        self.client.will_set(f"{self.base_topic}/availability", "offline", qos=1, retain=True)

        self.client.connect(host, port, keepalive=30)
        self.client.loop_start()

        # online markieren
        self.publish("availability", "online", retain=True)

    def publish(self, sub: str, payload, retain: Optional[bool] = None, qos: int = 0):
        if retain is None:
            retain = self.retain
        if isinstance(payload, (dict, list)):
            payload = json.dumps(payload, ensure_ascii=False)
        else:
            payload = str(payload)
        topic = f"{self.base_topic}/{sub.lstrip('/')}"
        self.client.publish(topic, payload, qos=qos, retain=retain)

    def close(self):
        try:
            self.publish("availability", "offline", retain=True, qos=1)
        except Exception:
            pass
        try:
            self.client.loop_stop()
        except Exception:
            pass
        try:
            self.client.disconnect()
        except Exception:
            pass


# =========================
# State
# =========================

@dataclass
class BaseState:
    ts: float
    uptime_s: float
    gnss_connected: bool
    lora_connected: bool

    # Survey-in
    svin_active: Optional[bool] = None
    svin_valid: Optional[bool] = None
    svin_dur_s: Optional[int] = None
    svin_mean_acc_m: Optional[float] = None
    svin_target_s: Optional[int] = None
    svin_target_acc_m: Optional[float] = None

    # Position
    fix: Optional[str] = None
    sats: Optional[int] = None
    lat: Optional[float] = None
    lon: Optional[float] = None
    h_m: Optional[float] = None
    hmsl_m: Optional[float] = None
    hacc_m: Optional[float] = None
    vacc_m: Optional[float] = None

    # RTCM
    rtcm_ok: int = 0
    rtcm_bad: int = 0
    rtcm_bytes: int = 0
    rtcm_types: Optional[List[int]] = None
    rtcm_last_ts: Optional[float] = None

    # LoRa
    lora_tx_bytes: int = 0


# =========================
# Main loop
# =========================

def open_serial(port: str, baud: int, timeout: float = 0.1) -> serial.Serial:
    return serial.Serial(port, baud, timeout=timeout)

def main():
    ap = argparse.ArgumentParser(description="Robuste RTK Base: ZED-F9P + optional LoRa TX + MQTT Status")
    ap.add_argument("--zed", "--port", dest="zed_port", default="/dev/ttyACM0")
    ap.add_argument("--zed-baud", dest="zed_baud", type=int, default=115200)

    ap.add_argument("--lora", dest="lora_port", default="/dev/ttyUSB0")
    ap.add_argument("--lora-baud", dest="lora_baud", type=int, default=57600)
    ap.add_argument("--lora-max", dest="lora_max", type=int, default=50)