"""Persistent storage for the 32-byte BLE payload-encryption key(s).

Each M5Stack generates its own key on first boot (per-device, persisted
to LittleFS). The host pulls the device's key during every profile
upload and stores it locally so subsequent BLE frames can be decrypted.

Storage:
    - ``config/.blekeyfile``        — legacy single-key file. The most
      recently uploaded key. Kept as a fallback for situations where a
      frame's device tag isn't in the per-MAC store.
    - ``config/.ble_keys.json``     — per-MAC key store, keyed by the
      eFuse base MAC (the same MAC that appears in each device's tag
      string "M5Stack|AA:BB:CC:DD:EE:FF"). Multiple ATOMS3s can coexist
      here so the user doesn't have to re-upload to switch devices.

Public API (back-compat preserved):
    load_key()                 -> bytes | None   (legacy file)
    save_key(key)              -> None           (writes legacy file)
    save_key_for_mac(mac, key) -> None           (writes per-MAC store)
    load_key_for_mac(mac)      -> bytes | None
    all_known_macs()           -> list[str]
"""

import json
import os
import threading
from utils.constants import APPDATA_DIR

KEY_LEN = 32
_KEY_FILE = os.path.join(APPDATA_DIR, ".blekeyfile")
_PER_MAC_FILE = os.path.join(APPDATA_DIR, ".ble_keys.json")

_lock = threading.Lock()


def key_path() -> str:
    return _KEY_FILE


# ---- Legacy single-key API ----

def load_key() -> bytes | None:
    try:
        with open(_KEY_FILE, "rb") as f:
            data = f.read()
    except OSError:
        return None
    return data if len(data) == KEY_LEN else None


def save_key(key: bytes) -> None:
    if len(key) != KEY_LEN:
        raise ValueError(f"BLE key must be {KEY_LEN} bytes, got {len(key)}")
    os.makedirs(APPDATA_DIR, exist_ok=True)

    # Detect a key change so we can wipe replay-protection counters: a new
    # key invalidates any captured ciphertext, so old counters carry no
    # protection value AND would block legitimate frames from a re-flashed
    # device until we manually reset.
    previous = load_key()

    tmp = _KEY_FILE + ".tmp"
    with open(tmp, "wb") as f:
        f.write(key)
    os.replace(tmp, _KEY_FILE)

    if previous != key:
        try:
            import ble_replay
            ble_replay.ReplayState().reset()
        except Exception:
            # Best-effort — replay reset failure isn't fatal (worst case,
            # the next BLE frame from the device gets rejected and the user
            # has to re-upload or wipe).
            pass


# ---- Per-MAC keystore ----

def _normalize_mac(mac: str) -> str:
    """Canonicalize MAC string: uppercase, colon-separated. Accepts the
    device tag "M5Stack|AA:BB:..." OR a bare MAC."""
    if mac.startswith("M5Stack|"):
        mac = mac[len("M5Stack|"):]
    return mac.upper()


def _load_per_mac_locked() -> dict:
    try:
        with open(_PER_MAC_FILE, "r", encoding="utf-8") as f:
            d = json.load(f)
    except (OSError, json.JSONDecodeError):
        return {}
    if not isinstance(d, dict):
        return {}
    return d


def _save_per_mac_locked(d: dict) -> None:
    os.makedirs(APPDATA_DIR, exist_ok=True)
    tmp = _PER_MAC_FILE + ".tmp"
    with open(tmp, "w", encoding="utf-8") as f:
        json.dump(d, f, indent=2)
    os.replace(tmp, _PER_MAC_FILE)


def save_key_for_mac(mac: str, key: bytes) -> None:
    """Persist ``key`` (32 raw bytes) under ``mac``. Replaces any
    existing key for that MAC and resets replay counters if the key
    actually changed."""
    if len(key) != KEY_LEN:
        raise ValueError(f"BLE key must be {KEY_LEN} bytes, got {len(key)}")
    mac = _normalize_mac(mac)
    with _lock:
        store = _load_per_mac_locked()
        prev_hex = store.get(mac)
        new_hex = key.hex()
        if prev_hex == new_hex:
            return
        store[mac] = new_hex
        _save_per_mac_locked(store)
    try:
        import ble_replay
        ble_replay.ReplayState().reset()
    except Exception:
        pass


def load_key_for_mac(mac: str) -> bytes | None:
    """Look up the key for ``mac`` (accepts bare MAC or full tag).
    Returns None if no key is stored for this device."""
    mac = _normalize_mac(mac)
    with _lock:
        store = _load_per_mac_locked()
    hex_key = store.get(mac)
    if not hex_key:
        return None
    try:
        b = bytes.fromhex(hex_key)
    except ValueError:
        return None
    return b if len(b) == KEY_LEN else None


def load_key_for_mac_or_default(mac: str) -> bytes | None:
    """Per-MAC lookup with the legacy single-key file as a fallback.
    Use this in the BLE layer when receiving a frame: try the right
    key for the device's MAC first, fall back to the legacy file for
    users who haven't re-uploaded since multi-device support landed."""
    k = load_key_for_mac(mac)
    if k is not None:
        return k
    return load_key()


def all_known_macs() -> list:
    with _lock:
        return sorted(_load_per_mac_locked().keys())