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

import argparse
import json
import os
import socket
import time
from typing import Any, Dict, Optional

import paho.mqtt.client as mqtt


def parse_args():
    ap = argparse.ArgumentParser(description="RTK Status -> MQTT (reads /run/rtk/state.json)")
    ap.add_argument("--state-file", default="/run/rtk/state.json")
    ap.add_argument("--mqtt-host", required=True)
    ap.add_argument("--mqtt-port", type=int, default=1883)
    ap.add_argument("--mqtt-user", default="")
    ap.add_argument("--mqtt-pass", default="")
    ap.add_argument("--topic-prefix", default="rtk/base")
    ap.add_argument("--interval", type=float, default=1.0)
    ap.add_argument("--retain", action="store_true", help="Retain MQTT publishes (default: false)")
    ap.add_argument("--client-id", default="rtk-status-base")
    ap.add_argument("--availability-topic", default="", help="Override availability topic (default: <prefix>/availability)")
    ap.add_argument("--debug", action="store_true")

    # Home Assistant MQTT Discovery
    ap.add_argument("--ha-discovery", action="store_true", help="Publish Home Assistant MQTT Discovery config")
    ap.add_argument("--ha-prefix", default="homeassistant", help="Discovery prefix (default: homeassistant)")
    ap.add_argument("--device-name", default="RTK Base", help="Device name shown in Home Assistant")
    ap.add_argument("--device-id", default="rtk_base", help="Stable device id for Home Assistant (no spaces)")
    return ap.parse_args()


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


def safe_read_json(path: str) -> Optional[Dict[str, Any]]:
    try:
        with open(path, "r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def flatten(prefix: str, obj: Any, out: Dict[str, Any]):
    """Flatten nested dicts into MQTT-friendly key paths."""
    if isinstance(obj, dict):
        for k, v in obj.items():
            key = f"{prefix}/{k}" if prefix else str(k)
            flatten(key, v, out)
    elif isinstance(obj, list):
        out[prefix] = json.dumps(obj, ensure_ascii=False)
    else:
        out[prefix] = obj


def _mqtt_client(client_id: str) -> mqtt.Client:
    # Robust gegen paho 1.x / 2.x Unterschiede
    return mqtt.Client(client_id=client_id, protocol=mqtt.MQTTv311)


def publish_discovery(client: mqtt.Client, args, availability_topic: str):
    """
    Publish Home Assistant MQTT Discovery configs.
    """

    dev = {
        "identifiers": [args.device_id],
        "name": args.device_name,
        "manufacturer": "u-blox / custom",
        "model": "ZED-F9P RTK",
        "sw_version": "rtk-status_mqtt",
    }

    base = args.topic_prefix.rstrip("/")
    ha = args.ha_prefix.rstrip("/")

    def pub_config(component: str, object_id: str, payload: Dict[str, Any]):
        topic = f"{ha}/{component}/{args.device_id}/{object_id}/config"
        client.publish(topic, json.dumps(payload, ensure_ascii=False), qos=0, retain=True)

    def common(name: str, state_topic: str):
        return {
            "name": name,
            "state_topic": state_topic,
            "availability_topic": availability_topic,
            "payload_available": "online",
            "payload_not_available": "offline",
            "device": dev,
        }

    def sensor_cfg(
        object_id: str,
        name: str,
        *,
        unit: Optional[str] = None,
        icon: Optional[str] = None,
        device_class: Optional[str] = None,
        state_class: Optional[str] = None,
        entity_category: Optional[str] = None,
    ):
        payload = {
            **common(name, f"{base}/{object_id}"),
            "unique_id": f"{args.device_id}_{object_id}",
        }
        if unit:
            payload["unit_of_measurement"] = unit
        if icon:
            payload["icon"] = icon
        if device_class:
            payload["device_class"] = device_class
        if state_class:
            payload["state_class"] = state_class
        if entity_category:
            payload["entity_category"] = entity_category
        pub_config("sensor", object_id, payload)

    def binary_sensor_cfg(
        object_id: str,
        name: str,
        *,
        icon: Optional[str] = None,
        device_class: Optional[str] = None,
        entity_category: Optional[str] = None,
    ):
        payload = {
            **common(name, f"{base}/{object_id}"),
            "unique_id": f"{args.device_id}_{object_id}",
            "payload_on": "true",
            "payload_off": "false",
        }
        if icon:
            payload["icon"] = icon
        if device_class:
            payload["device_class"] = device_class
        if entity_category:
            payload["entity_category"] = entity_category
        pub_config("binary_sensor", object_id, payload)

    # Existing / general sensors
    sensor_cfg("uptime_s", "RTK Uptime", unit="s", device_class="duration", icon="mdi:timer-outline")
    sensor_cfg("svin_meanAcc_m", "SVIN Mean Accuracy", unit="m", icon="mdi:crosshairs-gps", state_class="measurement")
    sensor_cfg("svin_state", "SVIN State", icon="mdi:satellite-variant")
    sensor_cfg("rtcm_last_s", "RTCM Age", unit="s", icon="mdi:timer-sand", state_class="measurement")
    sensor_cfg("rtcm_out_frames", "RTCM Frames Out", icon="mdi:counter", state_class="measurement")
    sensor_cfg("rtcm_out_bytes", "RTCM Bytes Out", unit="B", icon="mdi:database", state_class="measurement")
    sensor_cfg("lora_tx_bytes", "LoRa TX Bytes", unit="B", icon="mdi:transmission-tower", state_class="measurement")

    binary_sensor_cfg("zed_ok", "ZED OK", icon="mdi:satellite-uplink")
    binary_sensor_cfg("lora_ok", "LoRa OK", icon="mdi:radio-handheld")
    binary_sensor_cfg("state_ok", "State File OK", icon="mdi:file-check-outline")

    # New base position / quality / TMODE / SVIN sensors
    sensor_cfg("base_lat", "Base Latitude", icon="mdi:latitude")
    sensor_cfg("base_lon", "Base Longitude", icon="mdi:longitude")
    sensor_cfg("base_h_msl_m", "Base Height MSL", unit="m", icon="mdi:image-filter-hdr", state_class="measurement")
    sensor_cfg("base_h_ell_m", "Base Height Ellipsoid", unit="m", icon="mdi:image-filter-center-focus", state_class="measurement")
    sensor_cfg("base_geoid_sep_m", "Base Geoid Separation", unit="m", icon="mdi:terrain", state_class="measurement")
    sensor_cfg("base_fix_type", "Base Fix Type Code", icon="mdi:satellite-variant", state_class="measurement")
    sensor_cfg("base_fix_name", "Base Fix Type", icon="mdi:satellite-variant")
    sensor_cfg("base_sats", "Base Satellites", icon="mdi:satellite-uplink", state_class="measurement")
    sensor_cfg("base_hacc_m", "Base Horizontal Accuracy", unit="m", icon="mdi:crosshairs-gps", state_class="measurement")
    sensor_cfg("base_vacc_m", "Base Vertical Accuracy", unit="m", icon="mdi:arrow-expand-vertical", state_class="measurement")

    sensor_cfg("tmode3_mode", "Base TMODE3 Code", icon="mdi:cog-transfer", state_class="measurement")
    sensor_cfg("tmode3_name", "Base TMODE3", icon="mdi:cog-transfer")

    binary_sensor_cfg("svin_active", "Survey-In Active", icon="mdi:timer-sand")
    binary_sensor_cfg("svin_valid", "Survey-In Valid", icon="mdi:check-decagram")
    sensor_cfg("svin_dur_s", "Survey-In Duration", unit="s", icon="mdi:timer-outline", state_class="measurement")
    sensor_cfg("svin_meanAcc_m", "Survey-In Mean Accuracy", unit="m", icon="mdi:ruler", state_class="measurement")

    sensor_cfg("svin_target_dur_s", "Survey-In Target Duration", unit="s", icon="mdi:timer-cog-outline", state_class="measurement")
    sensor_cfg("svin_target_acc_m", "Survey-In Target Accuracy", unit="m", icon="mdi:target", state_class="measurement")


def main():
    args = parse_args()
    base = args.topic_prefix.rstrip("/")
    availability_topic = args.availability_topic.strip() or f"{base}/availability"

    host = socket.gethostname()

    client = _mqtt_client(args.client_id)

    connected = {"ok": False}

    def on_connect(client, userdata, flags, rc, *props):
        connected["ok"] = (rc == 0)
        if args.debug:
            print(f"[{now_iso()}] MQTT on_connect rc={rc}")

        client.publish(availability_topic, "online", retain=True)

        if args.ha_discovery:
            publish_discovery(client, args, availability_topic)

    def on_disconnect(client, userdata, rc, *props):
        connected["ok"] = False
        if args.debug:
            print(f"[{now_iso()}] MQTT on_disconnect rc={rc}")

    client.on_connect = on_connect
    client.on_disconnect = on_disconnect

    if args.mqtt_user:
        client.username_pw_set(args.mqtt_user, args.mqtt_pass)

    client.will_set(availability_topic, "offline", retain=True)

    client.connect(args.mqtt_host, args.mqtt_port, 30)
    client.loop_start()

    last_payload_hash = None

    while True:
        st = safe_read_json(args.state_file)

        flat: Dict[str, Any] = {}
        if st is not None:
            flatten("", st, flat)
            flat["host"] = host
            flat["state_ok"] = True
        else:
            flat = {
                "host": host,
                "state_ok": False,
                "publish_ts": now_iso(),
            }

        flat["publish_ts"] = now_iso()

        json_payload = json.dumps(
            flat if st is None else {**st, "host": host, "state_ok": True, "publish_ts": flat["publish_ts"]},
            ensure_ascii=False
        )

        h = hash(json_payload)
        if h != last_payload_hash:
            client.publish(f"{base}/json", json_payload, retain=args.retain)
            last_payload_hash = h

        for k, v in flat.items():
            topic = f"{base}/{k}"
            if isinstance(v, bool):
                payload = "true" if v else "false"
            elif v is None:
                payload = ""
            else:
                payload = str(v)
            client.publish(topic, payload, retain=args.retain)

        client.publish(availability_topic, "online", retain=True)

        time.sleep(max(0.2, args.interval))


if __name__ == "__main__":
    main()