"""Synchronous serial link to the STM32 debug protocol.

Provides blocking read/write of telemetry and parameters, suitable
for automated tuning scripts (not a TUI).  Mirrors the binary protocol
from code64/debug_console/protocol.py (kept in sync with the firmware's
debug_protocol.h — CRC-16, 78-byte telemetry, current parameter map).
"""

from __future__ import annotations

import struct
import time
from dataclasses import dataclass, field
from typing import Optional

import serial

# ── Protocol constants ───────────────────────────────────────────────

SYNC_BYTE = 0xAA

CMD_TELEMETRY = 0x01
CMD_PARAM_WRITE = 0x02
CMD_PARAM_WRITE_ACK = 0x03
CMD_PARAM_READ_ALL = 0x04
CMD_PARAM_VALUE = 0x05
CMD_PING = 0x10
CMD_PONG = 0x11
CMD_SHUTDOWN = 0x12          # turn off converter
CMD_RESET = 0x13             # system reset
CMD_TEST_50 = 0x14           # 50% duty test mode
CMD_RELAY_ON = 0x15          # latch input relay closed (bench test)
CMD_RELAY_OFF = 0x16         # latch input relay open  (bench test)
CMD_HOLD_CONVERTER = 0x17    # toggle "hold converter off" (boot guard + disarm trips)
CMD_TOGGLE_PRECHARGE = 0x18  # toggle the precharge FET (bench test)
CMD_ERROR_MSG = 0xE0

PTYPE_FLOAT = 0
PTYPE_UINT16 = 1
PTYPE_UINT8 = 2
PTYPE_INT32 = 3
PTYPE_INT16 = 4  # wire format = sign-extended int32, stored firmware-side as int16_t


# ── CRC-16/CCITT-FALSE (poly 0x1021, init 0xFFFF, no reflection) ──────
# Matches the STM32 hardware CRC unit configured in main.c MX_CRC_Init.

def crc16(data: bytes) -> int:
    crc = 0xFFFF
    for b in data:
        crc ^= b << 8
        for _ in range(8):
            if crc & 0x8000:
                crc = ((crc << 1) ^ 0x1021) & 0xFFFF
            else:
                crc = (crc << 1) & 0xFFFF
    return crc


# ── Telemetry ────────────────────────────────────────────────────────

@dataclass
class Telemetry:
    """Decoded telemetry packet from the STM32 (78-byte payload)."""
    vin: float = 0.0       # mV
    vout: float = 0.0      # mV
    iin: float = 0.0       # mA (negative = into converter)
    iout: float = 0.0      # mA
    vfly: float = 0.0      # mV
    etemp: float = 0.0     # °C (FET / external)
    btemp: float = 0.0     # °C (board)
    last_tmp: int = 0
    VREF: int = 0
    vfly_correction: int = 0
    cmp_outer: int = 0     # HRTIM Timer F CMP1xR (outer pair, T1/T4)
    vfly_integral: float = 0.0
    vfly_avg_debug: float = 0.0
    cc_output_f: float = 0.0
    mppt_iref: float = 0.0
    mppt_last_vin: float = 0.0
    mppt_last_iin: float = 0.0
    p_in: float = 0.0
    p_out: float = 0.0
    iout_slow: float = 0.0
    seq: int = 0
    cmp_inner: int = 0        # HRTIM Timer E CMP1xR (inner pair, T2/T3)
    vfly_ofs_applied: int = 0  # master-phase offset last written, signed ticks
    timestamp: float = field(default_factory=time.time)

    @property
    def vin_V(self) -> float:
        return self.vin / 1000.0

    @property
    def vout_V(self) -> float:
        return self.vout / 1000.0

    @property
    def iin_A(self) -> float:
        return self.iin / 1000.0

    @property
    def iout_A(self) -> float:
        return self.iout / 1000.0

    @property
    def power_in_W(self) -> float:
        return self.vin * (-self.iin) / 1e6

    @property
    def power_out_W(self) -> float:
        return self.vout * self.iout / 1e6

    @property
    def efficiency(self) -> float:
        p_in = self.power_in_W
        return (self.power_out_W / p_in * 100.0) if p_in > 0.1 else 0.0


_TELEM_FMT = "<7f hHhH 6f 3f BxH h"  # 78 bytes
_TELEM_SIZE = struct.calcsize(_TELEM_FMT)


def _decode_telemetry(payload: bytes) -> Optional[Telemetry]:
    if len(payload) < _TELEM_SIZE:
        return None
    v = struct.unpack(_TELEM_FMT, payload[:_TELEM_SIZE])
    return Telemetry(
        vin=v[0], vout=v[1], iin=v[2], iout=v[3], vfly=v[4], etemp=v[5], btemp=v[6],
        last_tmp=v[7], VREF=v[8], vfly_correction=v[9], cmp_outer=v[10],
        vfly_integral=v[11], vfly_avg_debug=v[12],
        cc_output_f=v[13], mppt_iref=v[14],
        mppt_last_vin=v[15], mppt_last_iin=v[16],
        p_in=v[17], p_out=v[18], iout_slow=v[19],
        seq=v[20], cmp_inner=v[21], vfly_ofs_applied=v[22],
    )


# ── Parameter definitions ────────────────────────────────────────────

@dataclass
class ParamDef:
    id: int
    name: str
    ptype: int
    group: str
    min_val: float = -1e9
    max_val: float = 1e9
    fmt: str = ".4f"


# Mirrors code64/debug_console/protocol.py PARAMS (firmware debug_protocol.c).
PARAMS = [
    # Compensator
    ParamDef(0x25, "VREF",             PTYPE_UINT16, "Compensator", 2340, 3500, ".0f"),
    # Vfly
    ParamDef(0x20, "vfly_kp",          PTYPE_FLOAT,  "Vfly", -10, 10, ".4f"),
    ParamDef(0x21, "vfly_ki",          PTYPE_FLOAT,  "Vfly", -10, 10, ".6f"),
    ParamDef(0x62, "vfly_kp_phase",    PTYPE_FLOAT,  "Vfly", -10, 10, ".4f"),   # mode 2: P gain, error -> phase
    ParamDef(0x63, "vfly_phase_clamp", PTYPE_UINT16, "Vfly", 0, 10000, ".0f"),  # mode 2: phase offset clamp
    ParamDef(0x22, "vfly_clamp",       PTYPE_UINT16, "Vfly", 0, 10000, ".0f"),
    ParamDef(0x23, "vfly_loop_trig",   PTYPE_UINT16, "Vfly", 1, 10000, ".0f"),
    ParamDef(0x24, "vfly_active",      PTYPE_UINT8,  "Vfly", 0, 3, ".0f"),      # 0=off 1=duty-asym PI 2=phase P 3=manual both
    ParamDef(0x26, "test_corr",        PTYPE_INT16,  "Vfly", -3000, 3000, ".0f"),  # mode-3 manual duty asymmetry
    ParamDef(0x27, "phase_ofs",        PTYPE_INT16,  "Vfly", -3000, 3000, ".0f"),  # master-phase: mode-3 manual, mode-2 readback
    # CC
    ParamDef(0x30, "cc_target",        PTYPE_FLOAT,  "CC", 0, 60000, ".0f"),
    ParamDef(0x31, "cc_gain",          PTYPE_FLOAT,  "CC", -1, 1, ".4f"),
    ParamDef(0x32, "cc_min_step",      PTYPE_FLOAT,  "CC", -1000, 0, ".1f"),
    ParamDef(0x33, "cc_max_step",      PTYPE_FLOAT,  "CC", 0, 1000, ".1f"),
    ParamDef(0x34, "cc_loop_trig",     PTYPE_UINT16, "CC", 1, 10000, ".0f"),
    ParamDef(0x35, "cc_active",        PTYPE_INT32,  "CC", 0, 1, ".0f"),
    # MPPT
    ParamDef(0x40, "mppt_step",        PTYPE_FLOAT,  "MPPT", 1, 200, ".0f"),
    ParamDef(0x41, "mppt_duty_min",    PTYPE_FLOAT,  "MPPT", 0, 6800, ".0f"),
    ParamDef(0x42, "mppt_duty_max",    PTYPE_FLOAT,  "MPPT", 0, 6800, ".0f"),
    ParamDef(0x44, "mppt_loop_trig",   PTYPE_UINT16, "MPPT", 1, 50000, ".0f"),
    ParamDef(0x45, "mppt_active",      PTYPE_INT32,  "MPPT", 0, 1, ".0f"),
    ParamDef(0x46, "cv_threshold",     PTYPE_FLOAT,  "MPPT", 20000, 30000, ".0f"),
    ParamDef(0x47, "cv_hysteresis",    PTYPE_FLOAT,  "MPPT", 0, 5000, ".0f"),
    ParamDef(0x48, "cc_threshold",     PTYPE_FLOAT,  "MPPT", 0, 55000, ".0f"),
    ParamDef(0x50, "cc_hysteresis",    PTYPE_FLOAT,  "MPPT", 0, 10000, ".0f"),
    # Deadtime (single static value, dt register units)
    ParamDef(0x60, "dt_normal",        PTYPE_UINT16, "Deadtime", 14, 200, ".0f"),
    # Manual fixed-duty mode (base duty in CMP ticks; D = override_duty/7158, 716..6442 = 10..90%)
    ParamDef(0x64, "override_duty",    PTYPE_UINT16, "Manual", 716, 6442, ".0f"),
    ParamDef(0x65, "manual_duty_en",   PTYPE_UINT8,  "Manual", 0, 1, ".0f"),
    # Closed-loop precharge PI (drives precharge FET PWM TIM3_CH1 to Vin/2)
    ParamDef(0x76, "precharge_kp",     PTYPE_FLOAT,  "Precharge", 0, 100, ".3f"),
    ParamDef(0x78, "precharge_ki",     PTYPE_FLOAT,  "Precharge", 0, 10, ".4f"),
    ParamDef(0x77, "precharge_reg_en", PTYPE_UINT8,  "Precharge", 0, 1, ".0f"),
    # Duty dither: delta-sigma the commanded duty between two out-of-band anchors (CMP ticks)
    ParamDef(0x70, "dither_en",        PTYPE_UINT8,  "Dither", 0, 1, ".0f"),
    ParamDef(0x71, "dither_band_lo",   PTYPE_UINT16, "Dither", 716, 6442, ".0f"),
    ParamDef(0x72, "dither_band_hi",   PTYPE_UINT16, "Dither", 716, 6442, ".0f"),
    ParamDef(0x73, "dither_anear",     PTYPE_UINT16, "Dither", 716, 6442, ".0f"),
    ParamDef(0x74, "dither_afar",      PTYPE_UINT16, "Dither", 716, 6442, ".0f"),
    ParamDef(0x75, "dither_dzero",     PTYPE_UINT16, "Dither", 716, 6442, ".0f"),
]

PARAM_BY_ID: dict[int, ParamDef] = {p.id: p for p in PARAMS}
PARAM_BY_NAME: dict[str, ParamDef] = {p.name: p for p in PARAMS}


# ── Frame building ───────────────────────────────────────────────────

def _build_frame(cmd: int, payload: bytes = b"") -> bytes:
    header = bytes([SYNC_BYTE, cmd, len(payload)])
    frame = header + payload
    crc = crc16(frame)
    return frame + bytes([(crc >> 8) & 0xFF, crc & 0xFF])  # big-endian: hi, lo


def _build_param_write(param_id: int, ptype: int, value) -> bytes:
    if ptype == PTYPE_FLOAT:
        val_bytes = struct.pack("<f", float(value))
    elif ptype == PTYPE_UINT16:
        val_bytes = struct.pack("<HH", int(value), 0)
    elif ptype == PTYPE_UINT8:
        val_bytes = struct.pack("<Bxxx", int(value))
    elif ptype == PTYPE_INT32:
        val_bytes = struct.pack("<i", int(value))
    elif ptype == PTYPE_INT16:
        val_bytes = struct.pack("<i", int(value))  # sign-extended 32-bit wire
    else:
        val_bytes = struct.pack("<I", int(value))
    payload = struct.pack("<BBxx", param_id, ptype) + val_bytes
    return _build_frame(CMD_PARAM_WRITE, payload)


def _decode_param_value(payload: bytes) -> Optional[tuple[int, float]]:
    if len(payload) < 8:
        return None
    param_id, ptype = payload[0], payload[1]
    vb = payload[4:8]
    if ptype == PTYPE_FLOAT:
        value = struct.unpack("<f", vb)[0]
    elif ptype == PTYPE_UINT16:
        value = float(struct.unpack("<H", vb[:2])[0])
    elif ptype == PTYPE_UINT8:
        value = float(vb[0])
    elif ptype == PTYPE_INT32:
        value = float(struct.unpack("<i", vb)[0])
    elif ptype == PTYPE_INT16:
        value = float(struct.unpack("<i", vb)[0])  # sign-extended 32-bit wire
    else:
        value = float(struct.unpack("<I", vb)[0])
    return (param_id, value)


# ── Frame parser state machine ───────────────────────────────────────

class _FrameParser:
    WAIT_SYNC = 0
    WAIT_CMD = 1
    WAIT_LEN = 2
    WAIT_PAYLOAD = 3
    WAIT_CRC_HI = 4
    WAIT_CRC_LO = 5

    def __init__(self):
        self.state = self.WAIT_SYNC
        self.cmd = 0
        self.length = 0
        self.buf = bytearray()
        self.payload = bytearray()
        self.idx = 0
        self.crc_hi = 0

    def feed(self, data: bytes):
        for b in data:
            if self.state == self.WAIT_SYNC:
                if b == SYNC_BYTE:
                    self.buf = bytearray([b])
                    self.state = self.WAIT_CMD
            elif self.state == self.WAIT_CMD:
                self.cmd = b
                self.buf.append(b)
                self.state = self.WAIT_LEN
            elif self.state == self.WAIT_LEN:
                self.length = b
                self.buf.append(b)
                self.payload = bytearray()
                self.idx = 0
                if b == 0:
                    self.state = self.WAIT_CRC_HI
                elif b > 128:
                    self.state = self.WAIT_SYNC
                else:
                    self.state = self.WAIT_PAYLOAD
            elif self.state == self.WAIT_PAYLOAD:
                self.payload.append(b)
                self.buf.append(b)
                self.idx += 1
                if self.idx >= self.length:
                    self.state = self.WAIT_CRC_HI
            elif self.state == self.WAIT_CRC_HI:
                self.crc_hi = b
                self.state = self.WAIT_CRC_LO
            elif self.state == self.WAIT_CRC_LO:
                received = (self.crc_hi << 8) | b
                expected = crc16(bytes(self.buf))
                self.state = self.WAIT_SYNC
                if received == expected:
                    yield (self.cmd, bytes(self.payload))


# ── STM32Link — synchronous serial interface ─────────────────────────

class STM32Link:
    """Blocking serial link to STM32 debug protocol.

    Usage::

        link = STM32Link("COM4")
        link.ping()
        t = link.read_telemetry()
        print(f"Vin={t.vin_V:.1f}V  Iout={t.iout_A:.1f}A  EFF={t.efficiency:.1f}%")
        link.write_param("dt_normal", 20)
        link.close()
    """

    def __init__(self, port: str, baudrate: int = 460800, timeout: float = 2.0):
        self.ser = serial.Serial(port, baudrate, timeout=timeout)
        self._parser = _FrameParser()
        self._param_cache: dict[int, float] = {}

    def close(self):
        if self.ser and self.ser.is_open:
            self.ser.close()

    def __enter__(self):
        return self

    def __exit__(self, *exc):
        self.close()

    # ── Low-level ────────────────────────────────────────────────────

    def _send(self, frame: bytes):
        self.ser.write(frame)

    def _recv_frames(self, timeout: float = 1.0) -> list[tuple[int, bytes]]:
        """Read available data and return decoded frames."""
        frames = []
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            data = self.ser.read(self.ser.in_waiting or 1)
            if data:
                for cmd, payload in self._parser.feed(data):
                    frames.append((cmd, payload))
            if frames:
                # Drain any remaining data
                time.sleep(0.02)
                data = self.ser.read(self.ser.in_waiting)
                if data:
                    for cmd, payload in self._parser.feed(data):
                        frames.append((cmd, payload))
                break
        return frames

    def _wait_for(self, target_cmd: int, timeout: float = 2.0) -> Optional[bytes]:
        """Wait for a specific command response, processing others."""
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                break
            data = self.ser.read(self.ser.in_waiting or 1)
            if data:
                for cmd, payload in self._parser.feed(data):
                    if cmd == target_cmd:
                        return payload
                    # Cache param values seen in passing
                    if cmd in (CMD_PARAM_VALUE, CMD_PARAM_WRITE_ACK):
                        result = _decode_param_value(payload)
                        if result:
                            self._param_cache[result[0]] = result[1]
                    # Cache telemetry too
                    if cmd == CMD_TELEMETRY:
                        self._last_telemetry = _decode_telemetry(payload)
        return None

    # ── Commands ─────────────────────────────────────────────────────

    def ping(self, timeout: float = 2.0) -> bool:
        """Send PING, return True if PONG received."""
        self._send(_build_frame(CMD_PING))
        return self._wait_for(CMD_PONG, timeout) is not None

    def shutdown(self):
        """Command the converter off."""
        self._send(_build_frame(CMD_SHUTDOWN))

    def reset(self):
        """Command a system reset."""
        self._send(_build_frame(CMD_RESET))

    def test_50(self):
        """Enter 50% duty test mode."""
        self._send(_build_frame(CMD_TEST_50))

    def relay_on(self):
        """Latch the input relay closed (bench test)."""
        self._send(_build_frame(CMD_RELAY_ON))

    def relay_off(self):
        """Latch the input relay open (bench test)."""
        self._send(_build_frame(CMD_RELAY_OFF))

    def hold_converter(self):
        """Toggle 'hold converter off' (boot guard + disarm trips)."""
        self._send(_build_frame(CMD_HOLD_CONVERTER))

    def toggle_precharge(self):
        """Toggle the precharge FET (bench test)."""
        self._send(_build_frame(CMD_TOGGLE_PRECHARGE))

    def read_telemetry(self, timeout: float = 2.0) -> Optional[Telemetry]:
        """Wait for next telemetry packet."""
        payload = self._wait_for(CMD_TELEMETRY, timeout)
        if payload:
            return _decode_telemetry(payload)
        return None

    def read_telemetry_avg(self, n: int = 10, timeout: float = 5.0) -> Optional[Telemetry]:
        """Read n telemetry packets and return the average."""
        samples: list[Telemetry] = []
        deadline = time.monotonic() + timeout
        while len(samples) < n and time.monotonic() < deadline:
            t = self.read_telemetry(timeout=deadline - time.monotonic())
            if t:
                samples.append(t)
        if not samples:
            return None
        # Average all analog float fields
        avg = Telemetry()
        for attr in ("vin", "vout", "iin", "iout", "vfly", "etemp", "btemp",
                      "vfly_integral", "vfly_avg_debug", "cc_output_f",
                      "mppt_iref", "mppt_last_vin", "mppt_last_iin",
                      "p_in", "p_out", "iout_slow"):
            setattr(avg, attr, sum(getattr(s, attr) for s in samples) / len(samples))
        avg.seq = samples[-1].seq
        return avg

    def request_all_params(self):
        """Request all parameter values from the STM32."""
        self._send(_build_frame(CMD_PARAM_READ_ALL))

    def read_all_params(self, timeout: float = 3.0) -> dict[str, float]:
        """Request and collect all parameter values."""
        self._param_cache.clear()
        self.request_all_params()
        deadline = time.monotonic() + timeout
        while time.monotonic() < deadline:
            data = self.ser.read(self.ser.in_waiting or 1)
            if data:
                for cmd, payload in self._parser.feed(data):
                    if cmd == CMD_PARAM_VALUE:
                        result = _decode_param_value(payload)
                        if result:
                            self._param_cache[result[0]] = result[1]
            time.sleep(0.05)
        # Convert to name->value
        return {
            PARAM_BY_ID[pid].name: val
            for pid, val in self._param_cache.items()
            if pid in PARAM_BY_ID
        }

    def write_param(self, name: str, value: float, wait_ack: bool = True) -> bool:
        """Write a parameter by name. Returns True if ACK received."""
        pdef = PARAM_BY_NAME.get(name)
        if not pdef:
            raise ValueError(f"Unknown parameter: {name!r}")
        if value < pdef.min_val or value > pdef.max_val:
            raise ValueError(
                f"{name}: {value} out of range [{pdef.min_val}, {pdef.max_val}]"
            )
        frame = _build_param_write(pdef.id, pdef.ptype, value)
        self._send(frame)
        if wait_ack:
            payload = self._wait_for(CMD_PARAM_WRITE_ACK, timeout=2.0)
            if payload:
                result = _decode_param_value(payload)
                if result:
                    self._param_cache[result[0]] = result[1]
                    return True
            return False
        return True

    def write_param_by_id(self, param_id: int, value: float) -> bool:
        """Write a parameter by ID."""
        pdef = PARAM_BY_ID.get(param_id)
        if not pdef:
            raise ValueError(f"Unknown param ID: 0x{param_id:02X}")
        return self.write_param(pdef.name, value)