"""Replay-protection state with per-session resilience.

Each side embeds (session_id, seq) in every frame. The session_id is a
random 64-bit value generated at boot/startup and stays constant for the
life of the process. The receiver tracks (last_session_id, max_seq) per
peer; a fresh session_id resets seq tracking gracefully (handles device
reflash, host wipe, etc.), while an unchanged session_id requires
strictly increasing seq (replay protection).

Tradeoff: an attacker who captures frames from an old session can replay
them AFTER the receiver has accepted a different session_id from the same
peer. For the user's threat model (passive sniffer in a controlled room,
no physical device access) this is acceptable. For stronger protection,
extend `device_state` to store a set of all ever-seen session_ids.
"""

import json
import os
import threading
import secrets
from pathlib import Path

from utils.constants import APPDATA_DIR

_REPLAY_STATE_PATH = Path(APPDATA_DIR) / ".ble_replay.json"


def _new_session_id() -> int:
    return secrets.randbits(64)


class ReplayState:
    """Per-MAC (session_id, max_seq) store + host's own session+counter."""

    def __init__(self):
        self._lock = threading.Lock()
        # Host's outgoing identity. `host_session_id` is regenerated on
        # every host restart so the device knows the host has reset and
        # accepts the (new session, fresh seq) gracefully.
        self._host_session_id: int = _new_session_id()
        self._host_send_seq: int = 0
        # Per-device tracking: {mac: {"session_id": int, "max_seq": int}}
        self._device_state: dict[str, dict] = {}
        self._load()

    def _load(self) -> None:
        try:
            with open(_REPLAY_STATE_PATH, "r", encoding="utf-8") as f:
                d = json.load(f)
            # We DELIBERATELY do NOT restore host_session_id from disk:
            # treating each Python process startup as a new session
            # automatically heals "host has stale state" cases.
            # We do persist host_send_seq within a session so a quick
            # crash-restart in the same process doesn't reuse seqs (but
            # since session_id changed, reuse is harmless anyway).
            raw_dev = d.get("device_state") or {}
            for mac, st in raw_dev.items():
                if isinstance(st, dict):
                    self._device_state[str(mac)] = {
                        "session_id": int(st.get("session_id", 0)),
                        "max_seq": int(st.get("max_seq", 0)),
                    }
            # Legacy migration: pre-session_id schema had `device_seen`.
            # We can't recover the old session_id (it didn't exist), so
            # we drop it — first frame from each device will be accepted
            # fresh under the new session_id.
        except (FileNotFoundError, json.JSONDecodeError, ValueError, TypeError):
            pass

    def _save_locked(self) -> None:
        try:
            _REPLAY_STATE_PATH.parent.mkdir(parents=True, exist_ok=True)
            tmp = _REPLAY_STATE_PATH.with_suffix(".tmp")
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump({
                    "host_send_seq": self._host_send_seq,
                    "device_state": self._device_state,
                }, f)
            os.replace(tmp, _REPLAY_STATE_PATH)
        except Exception:
            pass

    def host_session_id(self) -> int:
        return self._host_session_id

    def next_send_seq(self) -> int:
        """Reserve the next outgoing seq within the current host session."""
        with self._lock:
            self._host_send_seq += 1
            seq = self._host_send_seq
            self._save_locked()
            return seq

    def accept_received(self, mac: str, session_id: int, seq: int) -> tuple[bool, str]:
        """Validate an inbound (session_id, seq) from ``mac``.

        Returns (accepted, reason). ``reason`` is "fresh_session",
        "monotonic", "regressed_seq", or "stale_session".
        """
        if not isinstance(session_id, int) or not isinstance(seq, int):
            return False, "bad_types"
        with self._lock:
            st = self._device_state.get(mac)
            if st is None:
                # Never-seen device — accept and remember.
                self._device_state[mac] = {
                    "session_id": session_id,
                    "max_seq": seq,
                }
                self._save_locked()
                return True, "fresh_device"
            if st["session_id"] != session_id:
                # Different session: device rebooted (or was reflashed).
                # Accept fresh — this is the desync-recovery path.
                st["session_id"] = session_id
                st["max_seq"] = seq
                self._save_locked()
                return True, "fresh_session"
            # Same session: seq must be strictly monotonic.
            if seq <= st["max_seq"]:
                return False, "regressed_seq"
            st["max_seq"] = seq
            self._save_locked()
            return True, "monotonic"

    def reset(self) -> None:
        """Wipe all counters. Call when the AES key changes."""
        with self._lock:
            self._host_session_id = _new_session_id()
            self._host_send_seq = 0
            self._device_state = {}
            self._save_locked()