#!/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
from concurrent.futures import Future

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 mos_mode(charge_enabled: bool, discharge_enabled: bool) -> int:
    return (0 if charge_enabled else 1) | (0 if discharge_enabled else 2)


def cell_statistics(cells: list[float]) -> tuple[float, float, float] | None:
    if not cells:
        return None
    minimum = min(cells)
    maximum = max(cells)
    return minimum, maximum, round(maximum - minimum, 3)


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
    # Response checksums cover status, length, and payload, but not the command.
    expected = (-sum(frame[2:-3])) & 0xFFFF
    actual = int.from_bytes(frame[-3:-1], "big")
    return actual == 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.write_response = False
        self.buffer = bytearray()
        self.response_waiters: dict[int, asyncio.Future] = {}
        self.poll_lock = asyncio.Lock()
        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(value="")
        self.charge_enabled = tk.BooleanVar(value=False)
        self.discharge_enabled = tk.BooleanVar(value=False)
        self.mos_ready = False
        self.mos_busy = False
        self.rows: dict[str, tk.StringVar] = {}
        self.cell_rows: list[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", "Cell count"),
            ("cell_min", "Cell min"),
            ("cell_max", "Cell max"),
            ("cell_delta", "Cell delta"),
            ("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")

        cells = ttk.Frame(root)
        cells.grid(row=3, column=0, sticky="ew", pady=(8, 2))
        ttk.Label(cells, text="Cell voltages").grid(row=0, column=0, sticky="ne", padx=4, pady=2)
        self.cell_values = ttk.Frame(cells)
        self.cell_values.grid(row=0, column=1, sticky="w")

        controls = ttk.Frame(root)
        controls.grid(row=4, 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)
        self.charge_control = ttk.Checkbutton(
            controls,
            text="Charge enabled",
            variable=self.charge_enabled,
            command=self.set_mos,
        )
        self.charge_control.grid(row=0, column=2, padx=(12, 6))
        self.discharge_control = ttk.Checkbutton(
            controls,
            text="Discharge enabled",
            variable=self.discharge_enabled,
            command=self.set_mos,
        )
        self.discharge_control.grid(row=0, column=3, padx=6)
        self._update_mos_controls()

    def run(self, coro):
        future = asyncio.run_coroutine_threadsafe(coro, self.loop)
        future.add_done_callback(self._handle_future)
        return future

    def _handle_future(self, future: Future) -> None:
        try:
            future.result()
        except Exception as exc:
            self.root.after(0, lambda: self.status.set(f"BLE error: {exc}"))

    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.mos_ready = False
        self.mos_busy = False
        self._update_mos_controls()
        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()
        services = self.client.services
        if services is None:
            services = await self.client.get_services()
        write_char = services.get_characteristic(WRITE_UUID)
        notify_char = services.get_characteristic(NOTIFY_UUID)
        if write_char is None:
            raise RuntimeError(f"write characteristic not found: {WRITE_UUID}")
        if notify_char is None:
            raise RuntimeError(f"notify characteristic not found: {NOTIFY_UUID}")
        props = set(write_char.properties)
        self.write_response = "write-without-response" not in props and "write" in props
        await self.client.start_notify(NOTIFY_UUID, self._on_notify)
        await self._poll_once()

    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=self.write_response)

    def _on_notify(self, _sender, data: bytearray) -> None:
        self.buffer.extend(data)
        base_received = False
        for frame in split_frames(self.buffer):
            self.state.rx_log.append(frame.hex().upper())
            decode_response(frame, self.state)
            waiter = self.response_waiters.get(frame[1])
            if waiter and not waiter.done():
                waiter.set_result(frame)
            if validate_frame(frame) and frame[1] == 0x03 and frame[2] == 0:
                base_received = True
        self.root.after(0, lambda: self.status.set(f"Connected, TX {len(self.state.tx_log)}, RX {len(self.state.rx_log)}"))
        self.root.after(0, self._base_received if base_received else self._render_state)

    def _base_received(self) -> None:
        self.mos_ready = True
        self._render_state()
        self._update_mos_controls()

    def _update_mos_controls(self) -> None:
        state = "normal" if self.mos_ready and not self.mos_busy else "disabled"
        self.charge_control.configure(state=state)
        self.discharge_control.configure(state=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.charge_enabled.set(t.charge_mos)
        self.discharge_enabled.set(t.discharge_mos)
        self.rows["temperatures"].set(", ".join(f"{v:.1f} C" for v in t.temperatures) or "--")
        self.rows["cells"].set(str(len(t.cells)) if t.cells else "--")
        stats = cell_statistics(t.cells)
        self.rows["cell_min"].set(f"{stats[0]:.3f} V" if stats else "--")
        self.rows["cell_max"].set(f"{stats[1]:.3f} V" if stats else "--")
        self.rows["cell_delta"].set(f"{stats[2]:.3f} V" if stats else "--")
        self._render_cells(t.cells)
        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 _render_cells(self, cells: list[float]) -> None:
        if len(self.cell_rows) != len(cells):
            for widget in self.cell_values.winfo_children():
                widget.destroy()
            self.cell_rows = []
            for index in range(len(cells)):
                row = index // 4
                column = (index % 4) * 2
                value = tk.StringVar()
                self.cell_rows.append(value)
                ttk.Label(self.cell_values, text=f"Cell {index + 1}").grid(
                    row=row, column=column, sticky="e", padx=(4, 2), pady=2
                )
                ttk.Label(self.cell_values, textvariable=value, width=9).grid(
                    row=row, column=column + 1, sticky="w", padx=(0, 8), pady=2
                )
        for value, voltage in zip(self.cell_rows, cells):
            value.set(f"{voltage:.3f} V")

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

    async def _exchange(self, command: int, frame: bytes) -> bytes | None:
        waiter = self.loop.create_future()
        self.response_waiters[command] = waiter
        try:
            await self._send(frame)
            return await asyncio.wait_for(waiter, timeout=2.0)
        except TimeoutError:
            return None
        finally:
            if self.response_waiters.get(command) is waiter:
                del self.response_waiters[command]

    async def _request(self, command: int) -> bool:
        response = await self._exchange(command, read_frame(command))
        return bool(response and validate_frame(response) and response[2] == 0)

    async def _poll_once(self) -> None:
        if not self.client or not self.client.is_connected:
            self.root.after(0, lambda: self.status.set("Not connected"))
            return
        async with self.poll_lock:
            missing = []
            for index, command in enumerate((0x03, 0x04)):
                if index:
                    await asyncio.sleep(0.15)
                if not await self._request(command):
                    missing.append(f"0x{command:02X}")
            suffix = f", no response from {', '.join(missing)}" if missing else ""
            self.root.after(
                0,
                lambda: self.status.set(
                    f"Connected, TX {len(self.state.tx_log)}, RX {len(self.state.rx_log)}{suffix}"
                ),
            )

    def set_mos(self) -> None:
        if not self.mos_ready or self.mos_busy:
            self._render_state()
            return
        charge_enabled = self.charge_enabled.get()
        discharge_enabled = self.discharge_enabled.get()
        self.mos_busy = True
        self._update_mos_controls()
        self.status.set("Updating MOS state")
        self.run(self._set_mos(charge_enabled, discharge_enabled))

    async def _set_mos(self, charge_enabled: bool, discharge_enabled: bool) -> None:
        if not self.client or not self.client.is_connected:
            self.root.after(0, lambda: self._finish_mos_change("Not connected", ready=False))
            return
        async with self.poll_lock:
            mode = mos_mode(charge_enabled, discharge_enabled)
            frame = write_frame(0xE1, bytes([0x00, mode]))
            response = await self._exchange(0xE1, frame)
            accepted = bool(response and validate_frame(response) and response[2] == 0)
            refreshed = False
            if accepted:
                await asyncio.sleep(0.15)
                refreshed = await self._request(0x03)

        if not accepted:
            message = "MOS command was not acknowledged"
        elif not refreshed:
            message = "MOS updated; telemetry refresh timed out"
        else:
            message = f"MOS updated, TX {len(self.state.tx_log)}, RX {len(self.state.rx_log)}"
        self.root.after(0, lambda: self._finish_mos_change(message, ready=True))

    def _finish_mos_change(self, message: str, ready: bool) -> None:
        self.mos_busy = False
        self.mos_ready = ready
        self._render_state()
        self._update_mos_controls()
        self.status.set(message)

    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, self._show_disconnected)

    def _show_disconnected(self) -> None:
        self.mos_ready = False
        self.mos_busy = False
        self._update_mos_controls()
        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()