#!/usr/bin/env python3
"""Small desktop BLE client for Berger LFP / JBD-style battery telemetry."""

from __future__ import annotations

import asyncio
import threading
from dataclasses import dataclass, field

try:
    import tkinter as tk
    from tkinter import messagebox, ttk
except ModuleNotFoundError:  # Allows parser tests on minimal/headless Python.
    tk = None
    messagebox = None
    ttk = None

try:
    from bleak import BleakClient, BleakScanner
except ModuleNotFoundError:  # Allows parser tests before GUI dependencies are installed.
    BleakClient = None
    BleakScanner = None

SERVICE_UUID = "0000ff00-0000-1000-8000-00805f9b34fb"
NOTIFY_UUID = "0000ff01-0000-1000-8000-00805f9b34fb"
WRITE_UUID = "0000ff02-0000-1000-8000-00805f9b34fb"


def checksum(command: int, payload: bytes = b"") -> int:
    total = (command + len(payload) + sum(payload)) & 0xFFFF
    return ((total ^ 0xFFFF) + 1) & 0xFFFF


def build_frame(mode: int, command: int, payload: bytes = b"") -> bytes:
    check = checksum(command, payload)
    return bytes([0xDD, mode, command, len(payload), *payload, check >> 8, check & 0xFF, 0x77])


def read_frame(command: int) -> bytes:
    return build_frame(0xA5, command)


def write_frame(command: int, payload: bytes) -> bytes:
    return build_frame(0x5A, command, payload)


def validate_frame(frame: bytes) -> bool:
    if len(frame) < 7 or frame[0] != 0xDD or frame[-1] != 0x77:
        return False
    payload_len = frame[3]
    if len(frame) != payload_len + 7:
        return False
    expected = checksum(frame[1], frame[2:4] + frame[4 : 4 + payload_len])
    # Responses include mode/status at [1]/[2]; requests use command at [2].
    response_total = (frame[1] + frame[2] + frame[3] + sum(frame[4 : 4 + payload_len])) & 0xFFFF
    response_expected = ((response_total ^ 0xFFFF) + 1) & 0xFFFF
    actual = int.from_bytes(frame[-3:-1], "big")
    return actual in {expected, response_expected}


def split_frames(buffer: bytearray) -> list[bytes]:
    frames: list[bytes] = []
    while True:
        try:
            start = buffer.index(0xDD)
        except ValueError:
            buffer.clear()
            return frames
        if start:
            del buffer[:start]
        if len(buffer) < 7:
            return frames
        frame_len = buffer[3] + 7
        if len(buffer) < frame_len:
            return frames
        frame = bytes(buffer[:frame_len])
        del buffer[:frame_len]
        if frame[-1] == 0x77:
            frames.append(frame)


@dataclass
class Telemetry:
    voltage: float = 0.0
    current: float = 0.0
    power: float = 0.0
    soc: int = 0
    remaining_capacity: float = 0.0
    nominal_capacity: float = 0.0
    cycle_count: int = 0
    production_date: str = ""
    software_version: str = ""
    charge_mos: bool = False
    discharge_mos: bool = False
    balancing: bool = False
    protection: bool = False
    protection_bit: int | None = None
    temperatures: list[float] = field(default_factory=list)
    cells: list[float] = field(default_factory=list)
    time_label: str = "Est. time"
    time_value: str = "--"


@dataclass
class DeviceState:
    telemetry: Telemetry = field(default_factory=Telemetry)
    rx_log: list[str] = field(default_factory=list)
    tx_log: list[str] = field(default_factory=list)


def word(data: bytes, offset: int) -> int:
    return (data[offset] << 8) | data[offset + 1]


def signed_current(raw: int) -> float:
    if raw > 32768:
        raw -= 65536
    return round(raw / 100, 2)


def transform_time(hours: float) -> str:
    whole = int(hours)
    minutes = int((hours - whole) * 60)
    return f"{whole} h {minutes:02d} m"


def decode_response(frame: bytes, state: DeviceState) -> None:
    if not validate_frame(frame):
        return
    command = frame[1]
    status = frame[2]
    length = frame[3]
    payload = frame[4 : 4 + length]
    t = state.telemetry

    if command == 0x03 and status == 0 and length >= 23:
        t.voltage = round(word(payload, 0) / 100, 2)
        t.current = signed_current(word(payload, 2))
        t.power = round(t.voltage * t.current, 2)
        t.remaining_capacity = round(word(payload, 4) / 100, 2)
        t.nominal_capacity = round(word(payload, 6) / 100, 2)
        t.cycle_count = word(payload, 8)
        date_raw = word(payload, 10)
        day = date_raw & 0x1F
        month = (date_raw >> 5) & 0x0F
        year = 2000 + (date_raw >> 9)
        t.production_date = f"{year}-{month}-{day}"
        t.balancing = payload[12] != 0 or payload[13] != 0
        t.protection = payload[16] != 0 or payload[17] != 0
        if t.protection:
            bits = f"{payload[16]:08b}{payload[17]:08b}"[::-1]
            t.protection_bit = bits.find("1")
        version = str(payload[18])
        t.software_version = f"{version[0]}.{version[1]}" if len(version) > 1 else version
        t.soc = payload[19]
        mos = payload[20]
        t.charge_mos = mos in {1, 3}
        t.discharge_mos = mos in {2, 3}
        temp_count = payload[22] if length > 22 else 0
        temps = []
        for idx in range(temp_count):
            off = 23 + idx * 2
            if off + 1 < len(payload):
                temps.append(round((word(payload, off) - 2731) / 10, 1))
        t.temperatures = temps
        if t.current == 0:
            t.time_label = "Est. time"
            t.time_value = "--"
        elif t.current > 0 and t.nominal_capacity:
            t.time_label = "Time till full"
            t.time_value = transform_time((t.nominal_capacity - t.remaining_capacity) / t.current)
        elif t.current < 0:
            t.time_label = "Time till empty"
            t.time_value = transform_time(t.remaining_capacity / -t.current)

    elif command == 0x04 and status == 0:
        cells = []
        for off in range(0, len(payload) - 1, 2):
            cells.append(round(word(payload, off) / 1000, 3))
        t.cells = cells


def berger_match_reason(name: str | None, service_uuids: list[str]) -> str | None:
    normalized = {s.lower() for s in service_uuids}
    if SERVICE_UUID in normalized:
        return "service:ff00"
    if name and any(token in name.upper() for token in ("JBD", "BMS", "BERGER")):
        return "name"
    return None


class BergerApp:
    def __init__(self, root: tk.Tk) -> None:
        self.root = root
        self.root.title("Berger LFP BLE")
        self.state = DeviceState()
        self.devices = []
        self.client: BleakClient | None = None
        self.buffer = bytearray()
        self.loop = asyncio.new_event_loop()
        threading.Thread(target=self.loop.run_forever, daemon=True).start()

        self.status = tk.StringVar(value="Idle")
        self.scan_seconds = tk.IntVar(value=6)
        self.direct = tk.StringVar()
        self.rows: dict[str, tk.StringVar] = {}
        self._build_ui()

    def _build_ui(self) -> None:
        root = ttk.Frame(self.root, padding=10)
        root.grid(sticky="nsew")
        self.root.columnconfigure(0, weight=1)
        self.root.rowconfigure(0, weight=1)

        top = ttk.Frame(root)
        top.grid(row=0, column=0, sticky="ew")
        ttk.Button(top, text="Scan", command=self.scan).grid(row=0, column=0)
        ttk.Spinbox(top, from_=3, to=20, textvariable=self.scan_seconds, width=4).grid(row=0, column=1)
        ttk.Entry(top, textvariable=self.direct, width=32).grid(row=0, column=2, padx=6)
        ttk.Button(top, text="Connect", command=self.connect_selected).grid(row=0, column=3)
        ttk.Label(top, textvariable=self.status).grid(row=0, column=4, padx=8, sticky="w")

        self.listbox = tk.Listbox(root, height=8)
        self.listbox.grid(row=1, column=0, sticky="nsew", pady=8)

        data = ttk.Frame(root)
        data.grid(row=2, column=0, sticky="nsew")
        labels = [
            ("soc", "SOC"),
            ("voltage", "Voltage"),
            ("current", "Current"),
            ("power", "Power"),
            ("remaining_capacity", "Remaining Ah"),
            ("nominal_capacity", "Nominal Ah"),
            ("cycle_count", "Cycles"),
            ("mos", "MOS"),
            ("temperatures", "Temps"),
            ("cells", "Cells"),
            ("time", "Time"),
            ("production_date", "Production"),
            ("software_version", "Software"),
            ("status", "Status"),
        ]
        for idx, (key, label) in enumerate(labels):
            self.rows[key] = tk.StringVar(value="--")
            ttk.Label(data, text=label).grid(row=idx // 2, column=(idx % 2) * 2, sticky="e", padx=4, pady=2)
            ttk.Label(data, textvariable=self.rows[key], width=32).grid(row=idx // 2, column=(idx % 2) * 2 + 1, sticky="w")

        controls = ttk.Frame(root)
        controls.grid(row=3, column=0, sticky="ew", pady=8)
        ttk.Button(controls, text="Refresh", command=self.poll_once).grid(row=0, column=0)
        ttk.Button(controls, text="Disconnect", command=self.disconnect).grid(row=0, column=1, padx=6)

    def run(self, coro):
        return asyncio.run_coroutine_threadsafe(coro, self.loop)

    def scan(self) -> None:
        self.status.set("Scanning")
        self.run(self._scan())

    async def _scan(self) -> None:
        if BleakScanner is None:
            self.root.after(0, lambda: self.status.set("Install bleak first"))
            return
        devices = await BleakScanner.discover(timeout=self.scan_seconds.get(), return_adv=True)
        self.devices = []
        for device, adv in devices.values():
            services = list(adv.service_uuids or [])
            reason = berger_match_reason(device.name, services)
            self.devices.append((device, adv, reason))
        self.devices.sort(key=lambda item: (item[2] is None, -(item[1].rssi or -999)))
        self.root.after(0, self._render_devices)

    def _render_devices(self) -> None:
        self.listbox.delete(0, tk.END)
        for device, adv, reason in self.devices:
            mark = "*" if reason else " "
            self.listbox.insert(tk.END, f"{mark} {device.address}  {adv.rssi:>4}  {device.name or ''}  {reason or ''}")
        self.status.set(f"{len(self.devices)} devices")

    def connect_selected(self) -> None:
        address = self.direct.get().strip()
        if not address:
            selected = self.listbox.curselection()
            if not selected:
                messagebox.showinfo("Berger LFP BLE", "Select a device or enter an address.")
                return
            address = self.devices[selected[0]][0].address
        self.status.set("Connecting")
        self.run(self._connect(address))

    async def _connect(self, address: str) -> None:
        if BleakClient is None:
            self.root.after(0, lambda: self.status.set("Install bleak first"))
            return
        self.client = BleakClient(address)
        await self.client.connect()
        await self.client.start_notify(NOTIFY_UUID, self._on_notify)
        await self._send(read_frame(0x03))
        await self._send(read_frame(0x04))
        self.root.after(0, lambda: self.status.set("Connected"))

    async def _send(self, data: bytes) -> None:
        if not self.client:
            return
        self.state.tx_log.append(data.hex().upper())
        for off in range(0, len(data), 20):
            await self.client.write_gatt_char(WRITE_UUID, data[off : off + 20], response=False)

    def _on_notify(self, _sender, data: bytearray) -> None:
        self.buffer.extend(data)
        for frame in split_frames(self.buffer):
            self.state.rx_log.append(frame.hex().upper())
            decode_response(frame, self.state)
        self.root.after(0, self._render_state)

    def _render_state(self) -> None:
        t = self.state.telemetry
        self.rows["soc"].set(f"{t.soc} %")
        self.rows["voltage"].set(f"{t.voltage:.2f} V")
        self.rows["current"].set(f"{t.current:.2f} A")
        self.rows["power"].set(f"{t.power:.2f} W")
        self.rows["remaining_capacity"].set(f"{t.remaining_capacity:.2f} Ah")
        self.rows["nominal_capacity"].set(f"{t.nominal_capacity:.2f} Ah")
        self.rows["cycle_count"].set(str(t.cycle_count))
        self.rows["mos"].set(f"charge {'on' if t.charge_mos else 'off'}, discharge {'on' if t.discharge_mos else 'off'}")
        self.rows["temperatures"].set(", ".join(f"{v:.1f} C" for v in t.temperatures) or "--")
        self.rows["cells"].set(f"{len(t.cells)} cells, {min(t.cells):.3f}-{max(t.cells):.3f} V" if t.cells else "--")
        self.rows["time"].set(f"{t.time_label}: {t.time_value}")
        self.rows["production_date"].set(t.production_date or "--")
        self.rows["software_version"].set(t.software_version or "--")
        status = []
        if t.balancing:
            status.append("balancing")
        if t.protection:
            status.append(f"protection bit {t.protection_bit}")
        self.rows["status"].set(", ".join(status) or "normal")

    def poll_once(self) -> None:
        self.run(self._poll_once())

    async def _poll_once(self) -> None:
        await self._send(read_frame(0x03))
        await self._send(read_frame(0x04))

    def disconnect(self) -> None:
        self.run(self._disconnect())

    async def _disconnect(self) -> None:
        if self.client:
            await self.client.disconnect()
            self.client = None
        self.root.after(0, lambda: self.status.set("Disconnected"))


def main() -> None:
    if tk is None:
        raise SystemExit("tkinter is not installed. Install python3-tk to run the GUI.")
    root = tk.Tk()
    BergerApp(root)
    root.mainloop()


if __name__ == "__main__":
    main()