"""Stateless encoders/decoders for the Shining Mask BLE protocol.

This module contains *zero* I/O and *zero* BLE dependencies, so it can be unit
tested without hardware (and without ``bleak`` installed). Every ``cmd_*``
helper returns a ready-to-write, AES-encrypted 16-byte block for the command
characteristic. Upload data packets (sent unencrypted to the upload
characteristic) are produced by :func:`upload_packet`.

Command frame layout (before encryption), padded with 0x00 to 16 bytes::

    [ len ][ ASCII command name ][ args ... ][ 0x00 padding ... ]
      ^-- len = len(name) + len(args)

Notable divergences from mask-go (the field-tested Go reference), all chosen to
match the *real app traffic* documented in ble-protocol.md:

  * ``IMAG`` / ``ANIM`` use ``len = 5`` (name 4 + arg 1). mask-go hard-codes 6,
    which also works only because the extra byte lands in zero padding.
  * Background color command is ``BC`` (as captured from the app), not mask-go's
    ``BG`` typo.
  * Notification tokens are matched by *prefix*. (The doc *labels* the upload
    acks ``DATOK``/``REOKOK``, but its own hex decodes to ``DATSOK``/``REOKOK``,
    matching mask-go's observed ``DATSOK``/``REOK``. We match the real wire
    tokens ``DATSOK`` and ``REOK``.)
"""

from __future__ import annotations

from typing import Optional

import crypto
from constants import FRAME_SIZE, UPLOAD_MAX_DATA, UploadKind


# ---------------------------------------------------------------------------
# Command frame construction
# ---------------------------------------------------------------------------

def build_frame(name: str, args: bytes = b"") -> bytes:
    """Build a single *unencrypted* 16-byte command frame.

    ``len`` byte = number of (command name + argument) bytes, per the protocol
    spec. The remainder is zero padding (the mask uses AES-ECB, so padding is
    just zeros -- there is no IV/nonce).
    """
    name_b = name.encode("ascii")
    args = bytes(args)
    body = bytes([len(name_b) + len(args)]) + name_b + args
    if len(body) > FRAME_SIZE:
        raise ValueError(
            f"command {name!r} with {len(args)} arg bytes exceeds {FRAME_SIZE}-byte frame"
        )
    return body.ljust(FRAME_SIZE, b"\x00")


def encode_command(name: str, args: bytes = b"") -> bytes:
    """Build and AES-encrypt a command frame for the command characteristic."""
    return crypto.encrypt(build_frame(name, args))


def _u8(value: int, what: str) -> int:
    if not 0 <= value <= 255:
        raise ValueError(f"{what} must be 0..255, got {value}")
    return value


# ---------------------------------------------------------------------------
# Control commands (all return encrypted 16-byte blocks)
# ---------------------------------------------------------------------------

def cmd_brightness(level: int) -> bytes:
    """LIGHT: set brightness 0..255 (the doc recommends keeping it <= ~100)."""
    return encode_command("LIGHT", bytes([_u8(level, "brightness")]))


def cmd_image(image_id: int) -> bytes:
    """IMAG: show a built-in image by id."""
    return encode_command("IMAG", bytes([_u8(image_id, "image id")]))


def cmd_animation(anim_id: int) -> bytes:
    """ANIM: play a built-in animation by id."""
    return encode_command("ANIM", bytes([_u8(anim_id, "animation id")]))


def cmd_play_diy(image_id: int) -> bytes:
    """PLAY: show a single uploaded DIY image (count=1)."""
    return encode_command("PLAY", bytes([1, _u8(image_id, "DIY image id")]))


def cmd_play_diy_sequence(image_ids: "list[int]") -> bytes:
    """PLAY: play several uploaded DIY images in order (max 10)."""
    if not 1 <= len(image_ids) <= 10:
        raise ValueError("PLAY accepts between 1 and 10 DIY image ids")
    ids = bytes(_u8(i, "DIY image id") for i in image_ids)
    return encode_command("PLAY", bytes([len(image_ids)]) + ids)


def cmd_delete_diy(image_ids: "list[int]") -> bytes:
    """DELE: delete the given DIY images from the mask (max 10)."""
    if not 1 <= len(image_ids) <= 10:
        raise ValueError("DELE accepts between 1 and 10 DIY image ids")
    ids = bytes(_u8(i, "DIY image id") for i in image_ids)
    return encode_command("DELE", bytes([len(image_ids)]) + ids)


def cmd_check_diy_count() -> bytes:
    """CHEC: ask how many DIY images are stored (answer arrives on notify)."""
    return encode_command("CHEC")


def cmd_text_mode(mode: int) -> bytes:
    """MODE: scroll/animation mode for displayed text (see constants.TextMode)."""
    return encode_command("MODE", bytes([_u8(mode, "text mode")]))


def cmd_text_speed(speed: int) -> bytes:
    """SPEED: text scroll speed 0..255."""
    return encode_command("SPEED", bytes([_u8(speed, "text speed")]))


def cmd_text_color_effect(mode: int, enable: bool = True) -> bytes:
    """M: special text color effect / background image (see ble-protocol.md)."""
    return encode_command("M", bytes([1 if enable else 0, _u8(mode, "color effect")]))


def cmd_text_fg_color(r: int, g: int, b: int, enable: bool = True) -> bytes:
    """FC: foreground (text) color in RGB."""
    return encode_command(
        "FC", bytes([1 if enable else 0, _u8(r, "r"), _u8(g, "g"), _u8(b, "b")])
    )


def cmd_text_bg_color(r: int, g: int, b: int, enable: bool = True) -> bytes:
    """BC: background color in RGB (set to black to disable M's background images)."""
    return encode_command(
        "BC", bytes([1 if enable else 0, _u8(r, "r"), _u8(g, "g"), _u8(b, "b")])
    )


# ---------------------------------------------------------------------------
# Upload commands
# ---------------------------------------------------------------------------

def cmd_dats(total_len: int, second_field: int, kind: int) -> bytes:
    """DATS: announce an upload.

    For a text *bitmap* (kind=BITMAP), ``second_field`` is the bitmap length in
    bytes (so the mask knows where the appended color array starts). For a
    full-color *image* (kind=IMAGE), it is the destination image index. Both
    16-bit fields are big-endian.
    """
    if not 0 <= total_len <= 0xFFFF:
        raise ValueError("total_len out of 16-bit range")
    if not 0 <= second_field <= 0xFFFF:
        raise ValueError("second_field out of 16-bit range")
    args = total_len.to_bytes(2, "big") + second_field.to_bytes(2, "big") + bytes([_u8(kind, "kind")])
    return encode_command("DATS", args)


def cmd_datcp() -> bytes:
    """DATCP: finish/commit an upload. (mask-go omits the optional unix time.)"""
    return encode_command("DATCP")


def upload_packet(packet_index: int, chunk: bytes, pad_to: "Optional[int]" = None) -> bytes:
    """Build one *unencrypted* upload data packet for the upload characteristic.

    Layout: ``[ data_len+1 ][ packet_index ][ chunk... ]`` where ``data_len`` is
    ``len(chunk)``. The leading length byte counts the packet-index byte plus
    the data bytes. ``packet_index`` starts at 0 and increments per packet.

    For *bitmap* (text) uploads packets are NOT padded -- this matches mask-go,
    and the length byte already delimits the real content. For *image* uploads
    ble-protocol.md specifies fixed 100-byte packets, so pass ``pad_to=100``.
    """
    if not 0 <= packet_index <= 255:
        raise ValueError("packet_index must fit in one byte")
    if len(chunk) > UPLOAD_MAX_DATA:
        raise ValueError(f"chunk too large: {len(chunk)} > {UPLOAD_MAX_DATA}")
    packet = bytes([len(chunk) + 1, packet_index]) + bytes(chunk)
    if pad_to is not None and len(packet) < pad_to:
        packet = packet.ljust(pad_to, b"\x00")
    return packet


def iter_upload_packets(buffer: bytes, pad_to: "Optional[int]" = None):
    """Yield successive :func:`upload_packet` blocks covering ``buffer``."""
    index = 0
    for start in range(0, len(buffer), UPLOAD_MAX_DATA):
        yield upload_packet(index, buffer[start:start + UPLOAD_MAX_DATA], pad_to=pad_to)
        index += 1


# ---------------------------------------------------------------------------
# Audio visualization (optional / experimental -- per ble-protocol.md)
# ---------------------------------------------------------------------------

def cmd_audio_frame(mode: int, nibbles: "list[int]") -> bytes:
    """Build an encrypted audio-visualization frame for the audio characteristic.

    ``mode`` selects the visual style (0..4). ``nibbles`` is up to 28 row
    intensities (0x0..0xf), packed two per byte into 14 bytes. Frame layout:
    ``[ 0x0f ][ mode ][ 14 packed bytes ]`` = 16 bytes, then AES-encrypted.
    """
    if len(nibbles) > 28:
        raise ValueError("at most 28 nibbles (14 bytes) of visualization data")
    rows = list(nibbles) + [0] * (28 - len(nibbles))
    packed = bytes((rows[i] & 0x0F) << 4 | (rows[i + 1] & 0x0F) for i in range(0, 28, 2))
    frame = bytes([0x0F, _u8(mode, "audio mode")]) + packed
    return crypto.encrypt(frame)


# ---------------------------------------------------------------------------
# Notification parsing
# ---------------------------------------------------------------------------

def parse_notification(data: bytes) -> str:
    """Decrypt a notify-characteristic payload and return its ASCII token.

    The first decrypted byte is the token length; the token follows. Returns an
    empty string if the payload can't be interpreted.
    """
    if not data or len(data) % FRAME_SIZE != 0:
        return ""
    buf = crypto.decrypt(bytes(data))
    n = buf[0]
    if n == 0 or n > len(buf) - 1:
        return ""
    return buf[1:1 + n].decode("ascii", "ignore")


# Upload-handshake response tokens, matched by prefix. These are the real wire
# tokens (the doc's "DATOK" label is a typo -- its own hex decodes to "DATSOK").
RESP_DATS_OK = "DATSOK"   # ack of DATS (start of upload)
RESP_REC_OK = "REOK"      # ack of an upload packet (doc spells the full token "REOKOK")
RESP_DATCP_OK = "DATCPOK"  # ack of DATCP -> upload complete
RESP_PLAY_OK = "PLAYOK"
RESP_DELE_OK = "DELEOK"


def matches(token: str, expected: str) -> bool:
    """Prefix match for notification tokens (tolerant of spelling variants)."""
    return token.startswith(expected)