"""Encrypted-frame envelope shared by the variable-sync and live-keystroke
BLE channels.

Wire format (identical for both channels):
    byte tag_len               1 byte
    bytes tag                  ASCII "M5Stack|AA:BB:CC:DD:EE:FF"
    bytes nonce                12 bytes
    bytes ciphertext+gcm_tag   N + 16 bytes  (AES-256-GCM, AAD = tag bytes)

Each channel layers its own plaintext schema on top: the var-sync channel
uses JSON, the live-keystroke channel uses a compact binary protocol
defined in ble_live.py. Frames whose tag prefix isn't "M5Stack|" or whose
AEAD authentication fails are dropped silently by the parser.
"""

import os
import re

DEVICE_TAG_PREFIX = "M5Stack|"
MAC_RE = re.compile(r"^[0-9A-F]{2}(:[0-9A-F]{2}){5}$")


def build_frame(key: bytes, tag: str, plaintext: bytes) -> bytes:
    """Build an encrypted, MAC-tagged frame.

    `tag` is the device-tag string ("M5Stack|AA:BB:..."); it is bound to
    the ciphertext via GCM AAD and also written in plaintext at the head
    of the frame so the receiver can read it before decrypting.
    """
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM

    tag_bytes = tag.encode("ascii")
    if len(tag_bytes) > 255:
        raise ValueError("device tag too long")
    nonce = os.urandom(12)
    ct = AESGCM(key).encrypt(nonce, plaintext, tag_bytes)
    return bytes([len(tag_bytes)]) + tag_bytes + nonce + ct


def extract_tag(frame: bytes):
    """Pull the device tag (e.g. "M5Stack|AA:BB:CC:DD:EE:FF") out of
    the frame header WITHOUT decrypting. Returns the tag string or
    None if the frame is malformed.

    Used by the BLE layer to select the right per-MAC key before
    attempting AES-GCM authentication.
    """
    if len(frame) < 1 + 12 + 16:
        return None
    tag_len = frame[0]
    if tag_len == 0 or len(frame) < 1 + tag_len + 12 + 16:
        return None
    tag_bytes = bytes(frame[1:1 + tag_len])
    try:
        tag_str = tag_bytes.decode("ascii")
    except UnicodeDecodeError:
        return None
    if not tag_str.startswith(DEVICE_TAG_PREFIX):
        return None
    mac = tag_str[len(DEVICE_TAG_PREFIX):]
    if not MAC_RE.match(mac):
        return None
    return tag_str


def parse_frame(key: bytes, frame: bytes):
    """Return (tag_str, plaintext_bytes) on success, or None on any failure.

    Returns None silently for malformed / wrong-prefix / authentication-
    failed frames so callers can drop them without leaking timing info.
    """
    from cryptography.hazmat.primitives.ciphers.aead import AESGCM
    from cryptography.exceptions import InvalidTag

    tag_str = extract_tag(frame)
    if tag_str is None:
        return None
    tag_len = frame[0]
    tag_bytes = bytes(frame[1:1 + tag_len])
    nonce = bytes(frame[1 + tag_len:1 + tag_len + 12])
    ct = bytes(frame[1 + tag_len + 12:])
    try:
        pt = AESGCM(key).decrypt(nonce, ct, tag_bytes)
    except InvalidTag:
        return None
    except Exception:
        return None
    return tag_str, pt


def parse_frame_auto(frame: bytes):
    """Decode a frame using the right per-MAC key automatically.

    Reads the tag from the frame header, asks ble_keystore for the
    matching key (falling back to the legacy single-key file), then
    AES-GCM-decrypts. Returns (tag_str, plaintext_bytes) on success,
    or (tag_str, None) if a tag was readable but no matching key was
    available (so callers can surface a clear "unknown device" error
    instead of a silent auth fail), or None for fully malformed input.
    """
    import ble_keystore
    tag_str = extract_tag(frame)
    if tag_str is None:
        return None
    mac = tag_str[len(DEVICE_TAG_PREFIX):]
    key = ble_keystore.load_key_for_mac_or_default(mac)
    if key is None:
        return (tag_str, None)
    parsed = parse_frame(key, frame)
    if parsed is None:
        # Key existed but didn't authenticate — surface as a separate
        # signal (None for plaintext) so the live client can emit
        # KEY_MISMATCH for this specific device.
        return (tag_str, None)
    return parsed