"""EPS experiment bench: drive the EPS on its own, watch what it says back.

This is the bridge between replay and synthesis. Instead of replaying a
recording frame-for-frame, it transmits a chosen set of messages on a
fixed cycle - each ID at its own period, with a payload you control - so
you can answer "which messages does the EPS actually need to come up, and
what does it report while it's up?" by adding/removing entries rather than
re-recording.

Runs its own thread so timing doesn't depend on the UI's rerun cadence,
and tracks per-ID receive stats (count, rate, last payload, changed bytes)
for whatever the EPS transmits back.
"""
from __future__ import annotations

import csv
import sys
import threading
import time
from collections import deque
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "files"))
from decode_ptcan import CHECKSUMS, COUNTERS, ocsum  # noqa: E402


@dataclass
class TxEntry:
    """One message the bench transmits on a cycle.

    `sequence` holds the payloads this ID actually sent in the source
    capture, in order. Cycling through them reproduces the real alive
    counter and checksum progression byte-for-byte - which matters because
    a frozen payload reads as a stale sender, and not every checksum scheme
    on this bus is solved well enough to synthesise (0x1A0's only matches
    60-75% of the time, per files/PTCAN_protocol.md). Falls back to `data`
    when the sequence is empty.
    """
    arbitration_id: int
    data: bytes
    period_s: float
    enabled: bool = True
    sent: int = 0
    sequence: list[bytes] = field(default_factory=list)
    _next_due: float = 0.0
    _counter: int = 0
    _seq_i: int = 0

    def next_payload(self, refresh_counters: bool) -> bytes:
        if self.sequence:
            payload = self.sequence[self._seq_i % len(self.sequence)]
            self._seq_i += 1
            return payload
        self._counter += 1
        if refresh_counters:
            return apply_counter_and_checksum(self.arbitration_id, self.data, self._counter)
        return self.data


def apply_counter_and_checksum(arbitration_id: int, data: bytes, counter: int) -> bytes:
    """Refresh a frame's alive counter and checksum in place.

    Retransmitting a captured payload verbatim leaves its alive counter
    frozen, which every consumer of these messages treats as a stale/faulty
    sender - so a static transmit set gets ignored even though the exact
    same bytes worked during replay. Positions and the one's-complement
    checksum scheme come from files/PTCAN_protocol.md.
    """
    out = bytearray(data)

    pos = COUNTERS.get(arbitration_id)
    if pos is not None:
        byte_i, shift = pos
        if byte_i < len(out):
            # 4-bit counter that runs 0..14 and skips 15.
            out[byte_i] = (out[byte_i] & ~(0x0F << shift) & 0xFF) | ((counter % 15) << shift)

    cs = CHECKSUMS.get(arbitration_id)
    if cs is not None:
        idx, const = cs
        if idx < len(out):
            rest = [out[i] for i in range(len(out)) if i != idx]
            out[idx] = (ocsum(rest) + const) & 0xFF

    return bytes(out)


@dataclass
class RxInfo:
    """What we've seen back from the EPS for one arbitration ID."""
    count: int = 0
    first_seen: float = 0.0
    last_seen: float = 0.0
    last_data: bytes = b""
    changed_mask: int = 0          # bits that have ever differed between frames
    recent: deque = field(default_factory=lambda: deque(maxlen=32))

    @property
    def hz(self) -> float:
        if len(self.recent) < 2:
            return 0.0
        span = self.recent[-1] - self.recent[0]
        return (len(self.recent) - 1) / span if span > 0 else 0.0


class EpsBench:
    def __init__(self, adapter, log_path: Optional[Path] = None):
        self.adapter = adapter
        self.tx: dict[int, TxEntry] = {}
        self.rx: dict[int, RxInfo] = {}
        self.refresh_counters = True   # keep alive counters/checksums live
        self._lock = threading.Lock()
        self._stop = threading.Event()
        self._thread: Optional[threading.Thread] = None
        self._t0 = 0.0
        self._log_file = None
        self._log_writer = None
        if log_path is not None:
            self._log_file = open(log_path, "w", newline="")
            self._log_writer = csv.writer(self._log_file)
            self._log_writer.writerow(["t", "wall", "dir", "arbitration_id", "dlc", "data_hex"])

    # ---- transmit set -----------------------------------------------------

    def set_entries(self, entries: list[TxEntry]) -> None:
        with self._lock:
            now = time.monotonic()
            for e in entries:
                e._next_due = now
            self.tx = {e.arbitration_id: e for e in entries}

    def set_enabled(self, arbitration_id: int, enabled: bool) -> None:
        with self._lock:
            if arbitration_id in self.tx:
                self.tx[arbitration_id].enabled = enabled

    def set_payload(self, arbitration_id: int, data: bytes) -> None:
        with self._lock:
            if arbitration_id in self.tx:
                self.tx[arbitration_id].data = data

    def enable_all(self, enabled: bool) -> None:
        with self._lock:
            for e in self.tx.values():
                e.enabled = enabled

    # ---- lifecycle --------------------------------------------------------

    def is_running(self) -> bool:
        return bool(self._thread and self._thread.is_alive())

    def start(self) -> None:
        if self.is_running():
            return
        self._stop.clear()
        self._t0 = time.monotonic()
        self._thread = threading.Thread(target=self._run, daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        if self._thread:
            self._thread.join(timeout=2)
        if self._log_file is not None:
            self._log_file.close()
            self._log_file = None
            self._log_writer = None  # a restart must not write to the closed file

    def reset_rx(self) -> None:
        with self._lock:
            self.rx.clear()

    def actual_rate(self) -> Optional[float]:
        """Frames/s actually achieved so far, or None before transmitting."""
        if not self._t0:
            return None
        elapsed = time.monotonic() - self._t0
        if elapsed <= 0:
            return None
        with self._lock:
            total = sum(e.sent for e in self.tx.values())
        return total / elapsed if total else None

    def run_startup_sequence(self, monitor, pre_bus_s: float = 1.0,
                             settle_s: float = 2.0) -> None:
        """Bring the EPS up in the order a real car does it.

        Order matters: the module wants a populated bus already talking when
        its 12V enable arrives, rather than waking into silence and latching
        a fault. So: start transmitting, let the bus look "alive" for a
        moment, then apply the enable, then hold while it initialises.
        Counters keep advancing throughout, which is the point.
        """
        self.reset_rx()
        if not self.is_running():
            self.start()
        time.sleep(pre_bus_s)
        monitor.set_output_verified(0, True)
        time.sleep(settle_s)

    def snapshot(self):
        with self._lock:
            return dict(self.tx), dict(self.rx)

    # ---- worker -----------------------------------------------------------

    def _run(self) -> None:
        while not self._stop.is_set():
            now = time.monotonic()
            with self._lock:
                due = [e for e in self.tx.values() if e.enabled and e._next_due <= now]
                payloads = []
                for e in due:
                    e._next_due = now + e.period_s
                    e.sent += 1
                    payloads.append(e.next_payload(self.refresh_counters))
            if due:
                self.adapter.send_frames([(e.arbitration_id, d, False) for e, d in zip(due, payloads)])
                if self._log_writer is not None:
                    wall = time.time()
                    for e, data in zip(due, payloads):
                        self._log_writer.writerow(
                            [round(now - self._t0, 4), round(wall, 4), "tx", f"{e.arbitration_id:X}",
                             len(data), data.hex().upper()]
                        )

            # Drain whatever the EPS sent back, but never past the next due
            # time - falling behind here is what makes messages arrive stale.
            with self._lock:
                deadline = min((e._next_due for e in self.tx.values() if e.enabled), default=now + 0.02)
            while time.monotonic() < deadline:
                frame = self.adapter.read_frame(timeout=0.001)
                if frame is None:
                    break
                self._record_rx(frame)
            slack = deadline - time.monotonic()
            if slack > 0:
                time.sleep(min(slack, 0.005))

    def _record_rx(self, frame) -> None:
        t = time.monotonic() - self._t0
        with self._lock:
            info = self.rx.get(frame.arbitration_id)
            if info is None:
                info = RxInfo(first_seen=t)
                self.rx[frame.arbitration_id] = info
            if info.last_data and len(info.last_data) == len(frame.data):
                for i, (a, b) in enumerate(zip(info.last_data, frame.data)):
                    if a != b:
                        info.changed_mask |= 1 << i
            info.count += 1
            info.last_seen = t
            info.last_data = frame.data
            info.recent.append(t)
        if self._log_writer is not None:
            self._log_writer.writerow(
                [round(t, 4), round(time.time(), 4), "rx", f"{frame.arbitration_id:X}",
                 len(frame.data), frame.data.hex().upper()]
            )


def entries_from_capture(path: Path, ids: Optional[set[int]] = None,
                         min_count: int = 2, max_sequence: int = 400) -> list[TxEntry]:
    """Build a transmit set from a capture: one entry per ID, using that ID's
    median period and the payload sequence it actually sent.

    Keeping the sequence (rather than just the last payload) is what makes a
    synthetic transmit set acceptable to the EPS - see TxEntry.
    """
    times: dict[int, list[float]] = {}
    payloads: dict[int, list[bytes]] = {}
    with open(path, newline="") as f:
        rows = sorted(csv.DictReader(f), key=lambda r: int(r["timestamp_ms"]))
        if not rows:
            return []
        t0 = int(rows[0]["timestamp_ms"])
        for row in rows:
            aid = int(row["arbitration_id"], 16)
            if ids is not None and aid not in ids:
                continue
            times.setdefault(aid, []).append((int(row["timestamp_ms"]) - t0) / 1000)
            seq = payloads.setdefault(aid, [])
            if len(seq) < max_sequence:
                seq.append(bytes.fromhex(row["data_hex"]))

    entries = []
    for aid, ts in times.items():
        if len(ts) < min_count:
            continue
        gaps = sorted(b - a for a, b in zip(ts, ts[1:]) if 0 < b - a < 5)
        period = gaps[len(gaps) // 2] if gaps else 1.0
        seq = payloads[aid]
        entries.append(TxEntry(arbitration_id=aid, data=seq[-1],
                               period_s=round(period, 3), sequence=seq))
    entries.sort(key=lambda e: e.arbitration_id)
    return entries


# Candidate subsets, smallest first. Narrowing matters for two reasons: it
# answers "what does the EPS actually need", and it cuts the frame rate -
# the full 69-ID set needs ~1330 frames/s, which the PC-side bench can't
# always sustain while the UI is also running, and late frames look exactly
# like a faulty sender to the EPS.
PRESETS: dict[str, set[int]] = {
    # Steering, road/wheel speed, engine state, terminal status: the inputs a
    # speed-sensitive power steering controller plausibly can't work without.
    "Minimal candidate": {
        0x0C4,   # steering angle + rate (SZL)
        0x1A0,   # road speed (DSC)
        0x0CE,   # wheel speeds (DSC)
        0x0AA,   # engine speed (DME)
        0x0A9,   # system voltage (DME)
        0x0A8,   # torque (DME)
        0x130,   # terminal status (CAS)
        0x19E,   # DSC status
        0x0B6,   # DSC counter/heartbeat
    },
    # Adds the rest of the fast chassis/powertrain traffic that was present
    # before the EPS woke in the reference capture.
    "Core chassis + powertrain": {
        0x0C4, 0x1A0, 0x0CE, 0x0AA, 0x0A9, 0x0A8, 0x130, 0x19E, 0x0B6,
        0x0C8, 0x1A6, 0x2B2, 0x1B6, 0x194, 0x1E1, 0x1D0,
    },
}


def required_rate(entries: list[TxEntry]) -> float:
    """Frames per second a transmit set demands, all entries enabled."""
    return sum(1.0 / e.period_s for e in entries if e.period_s > 0)