"""
kamoji_ble.py - talk to a Kamoji over Bluetooth Low Energy, and show your own photo or GIF on it.

    pip install bleak pillow

    python kamoji_ble.py info                      # firmware version, asset count, nickname
    python kamoji_ble.py photo cat.jpg --seconds 30
    python kamoji_ble.py gif dance.gif --seconds 60
    python kamoji_ble.py stop                      # end what is on screen now
    python kamoji_ble.py forget                    # Kamoji stops trusting this computer

The first time, hold your laptop close to Kamoji (in the car): Kamoji only trusts a new device
that is close by. The key that makes this laptop trusted is kept in ~/.kamoji_ble_key.

Protocol: https://docs.kamoji.bot/docs/modding/bluetooth-protocol
"""
import argparse
import asyncio
import hashlib
import hmac
import io
import os
import secrets
import struct
import sys
import time
from pathlib import Path

from bleak import BleakClient, BleakScanner
from PIL import Image, ImageSequence

SERVICE = "12345678-1234-5678-1234-56789abcdef0"
CMD = "12345678-1234-5678-1234-56789abcdef1"  # write: commands
TX = "12345678-1234-5678-1234-56789abcdef3"   # notify: answers

SIDE = 240                 # the Kamoji's screen, in pixels
MAX_GIF_BYTES = 250 * 1024
KEY_FILE = Path.home() / ".kamoji_ble_key"


# ---------------------------------------------------------------- connecting


async def find_kamoji(timeout=10.0):
    """The first Kamoji in range (it advertises as "KMJ-XXXX")."""
    device = await BleakScanner.find_device_by_filter(
        lambda d, ad: (d.name or "").startswith(("KMJ-", "KM-"))
        or SERVICE in (ad.service_uuids or []),
        timeout=timeout,
    )
    if device is None:
        sys.exit("No Kamoji found. Is it on, and not connected to the app?")
    return device


class Kamoji:
    """One connection: send(command) and wait_for(prefix) on the answers."""

    def __init__(self, client):
        self.client = client
        self.messages = asyncio.Queue()

    def _on_notify(self, _, data: bytearray):
        self.messages.put_nowait(data.decode(errors="replace").strip())

    async def start(self):
        await self.client.start_notify(TX, self._on_notify)
        await self.authorize()
        # Keep the link fast for the next 8 s (it slows down when idle).
        await self.send("LINK_FAST")

    async def authorize(self):
        """Prove this laptop is trusted (or become trusted, when close by). Older firmware has no
        authorization and answers UNKNOWN_CMD:AUTH_HELLO."""
        key = load_key()
        key_id = hashlib.sha256(key).digest()[:8].hex()
        await self.send(f"AUTH_HELLO:{key_id}")
        msg = await self.wait_for("AUTH:", "UNKNOWN_CMD:AUTH_HELLO", allow_unknown=True)
        if msg is None or msg.startswith("UNKNOWN_CMD:"):
            return  # firmware without authorization (or older, silent firmware)
        if msg.startswith("AUTH:CHALLENGE:"):
            nonce = bytes.fromhex(msg.split(":")[2])
            proof = hmac.new(key, nonce, hashlib.sha256).hexdigest()
            await self.send(f"AUTH_PROOF:{proof}")
        elif msg.startswith("AUTH:ENROLL:"):
            await self.send(f"AUTH_ENROLL:{key.hex()}")
        msg = msg if msg.startswith("AUTH:ERR:") else await self.wait_for("AUTH:OK:", "AUTH:ERR:")
        if msg and msg.startswith("AUTH:OK:"):
            if msg.split(":")[2] == "1":
                print("Kamoji now trusts this computer.")
            return
        if msg and msg.startswith("AUTH:ERR:FAR:"):
            sys.exit("Kamoji does not know this computer yet and it is too far away: "
                     "bring it close to Kamoji (in the car) and try again.")
        sys.exit(f"Kamoji did not accept this computer: {msg}")

    async def send(self, command: str):
        await self.client.write_gatt_char(CMD, command.encode(), response=True)

    async def wait_for(self, *prefixes, timeout=5.0, allow_unknown=False):
        """The next answer starting with one of `prefixes`, or None after `timeout` s."""
        deadline = time.monotonic() + timeout
        while (left := deadline - time.monotonic()) > 0:
            try:
                msg = await asyncio.wait_for(self.messages.get(), left)
            except asyncio.TimeoutError:
                break
            if msg.startswith("AUTH:REQUIRED:"):
                sys.exit(f"Kamoji refused {msg[14:]}: this computer is not authorized.")
            if msg.startswith("UNKNOWN_CMD:") and not allow_unknown:
                sys.exit(f"This Kamoji's firmware does not know {msg[12:]}: update it.")
            if msg.startswith(prefixes):
                return msg
        return None


def load_key() -> bytes:
    """This computer's 32-byte key, created on first use. Keep it private: whoever has it can
    control every Kamoji that trusts this computer."""
    if KEY_FILE.exists():
        key = bytes.fromhex(KEY_FILE.read_text().strip())
        if len(key) == 32:
            return key
    key = secrets.token_bytes(32)
    KEY_FILE.write_text(key.hex())
    os.chmod(KEY_FILE, 0o600)
    return key


# ---------------------------------------------------------------- preparing


def prepare_photo(path: str, quality=75) -> bytes:
    """Centre-crop, 240x240, baseline JPEG (the Kamoji's decoder reads baseline only)."""
    img = Image.open(path).convert("RGB")
    side = min(img.size)
    left, top = (img.width - side) // 2, (img.height - side) // 2
    img = img.crop((left, top, left + side, top + side)).resize((SIDE, SIDE), Image.LANCZOS)
    out = io.BytesIO()
    img.save(out, "JPEG", quality=quality, progressive=False, optimize=True)
    return out.getvalue()


def prepare_gif(path: str) -> bytes:
    """Centre-crop every frame to 240x240, then fewer frames / colours until <= 250 KB."""
    src = Image.open(path)
    frames, delays = [], []
    for frame in ImageSequence.Iterator(src):
        rgb = frame.convert("RGB")
        side = min(rgb.size)
        left, top = (rgb.width - side) // 2, (rgb.height - side) // 2
        frames.append(rgb.crop((left, top, left + side, top + side)).resize((SIDE, SIDE), Image.LANCZOS))
        delays.append(max(20, frame.info.get("duration", 100)))

    for keep_every, colors in [(1, 128), (2, 128), (2, 64), (3, 48), (4, 32)]:
        picked = frames[::keep_every]
        # Each kept frame stands in for the ones dropped after it.
        picked_delays = [sum(delays[i:i + keep_every]) for i in range(0, len(frames), keep_every)]
        palette = picked[0].quantize(colors)
        indexed = [f.quantize(palette=palette) for f in picked]
        out = io.BytesIO()
        indexed[0].save(out, "GIF", save_all=True, append_images=indexed[1:],
                        duration=picked_delays, loop=0, optimize=True, disposal=1)
        data = out.getvalue()
        print(f"  {len(picked)} frames, {colors} colours: {len(data) // 1024} KB")
        if len(data) <= MAX_GIF_BYTES:
            return data
    sys.exit("This GIF stays over 250 KB: try a shorter or simpler one.")


# ---------------------------------------------------------------- sending


async def show(kamoji: Kamoji, data: bytes, seconds: int):
    """IMG_BEGIN, the data in IMG_DATA chunks (acked every N bytes), then wait for IMG:OK."""
    t0 = time.monotonic()
    await kamoji.send(f"IMG_BEGIN:{len(data)}:{seconds}")
    ready = await kamoji.wait_for("IMG:READY:", "IMG:ERR:")
    if ready is None:
        sys.exit("No answer to IMG_BEGIN: this firmware has no photo support yet.")
    if ready.startswith("IMG:ERR:"):
        sys.exit(f"Kamoji refused it: {ready}")
    ack_every = int(ready.split(":")[2])

    # Each write: "IMG_DATA:" + offset (uint32, little-endian) + bytes, filling the MTU.
    chunk = max(20, kamoji.client.mtu_size - 3) - 13
    window = 2 * ack_every       # how far ahead of the last ack we may send
    acked, offset = 0, 0
    while acked < len(data):
        while offset < len(data) and offset < acked + window:
            part = data[offset:offset + chunk]
            frame = b"IMG_DATA:" + struct.pack("<I", offset) + part
            await kamoji.client.write_gatt_char(CMD, frame, response=False)
            offset += len(part)
        msg = await kamoji.wait_for("IMG:ACK:", "IMG:RESEND:", "IMG:ERR:", timeout=3)
        if msg is None:
            offset = acked                      # no ack: resend what was not acked
        elif msg.startswith("IMG:ERR:"):
            sys.exit(f"Transfer failed: {msg}")
        else:
            acked = max(acked, int(msg.split(":")[2]))
            if msg.startswith("IMG:RESEND:"):
                offset = int(msg.split(":")[2])  # a write was lost: resend from there
        print(f"\r  sent {acked * 100 // len(data)}%", end="", flush=True)
    print()

    shown = await kamoji.wait_for("IMG:OK", "IMG:ERR:", timeout=8)
    if shown != "IMG:OK":
        sys.exit(f"Kamoji could not show it: {shown}")
    print(f"On screen {time.monotonic() - t0:.1f} s after IMG_BEGIN, for {seconds} s.")


# ---------------------------------------------------------------- main


async def main():
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("action", choices=["info", "photo", "gif", "stop", "forget"])
    parser.add_argument("file", nargs="?")
    parser.add_argument("--seconds", type=int, default=30, help="how long it stays up (1-600)")
    args = parser.parse_args()
    if args.action in ("photo", "gif") and not args.file:
        parser.error(f"{args.action} needs a file")

    data = None
    if args.action == "photo":
        data = prepare_photo(args.file)
        print(f"Photo: {len(data) // 1024} KB")
    elif args.action == "gif":
        data = prepare_gif(args.file)

    device = await find_kamoji()
    print(f"Connecting to {device.name}...")
    async with BleakClient(device) as client:
        kamoji = Kamoji(client)
        if args.action == "forget":
            # No authorization needed: the token proves this computer's key.
            await client.start_notify(TX, kamoji._on_notify)
            key = load_key()
            key_id = hashlib.sha256(key).digest()[:8].hex()
            token = hmac.new(key, b"forget", hashlib.sha256).hexdigest()
            await kamoji.send(f"AUTH_FORGET:{key_id}:{token}")
            print(await kamoji.wait_for("AUTH:FORGOT:", "AUTH:ERR:", "UNKNOWN_CMD:", allow_unknown=True))
            return
        await kamoji.start()
        if args.action == "info":
            await kamoji.send("GET_INFO")
            for _ in range(4):  # FW:, AT:, AMD5:, NICK:
                msg = await kamoji.wait_for("FW:", "AT:", "AMD5:", "NICK:")
                if msg:
                    print(" ", msg)
        elif args.action == "stop":
            await kamoji.send("IMG_STOP")
            print("Stopped.")
        else:
            await show(kamoji, data, args.seconds)


if __name__ == "__main__":
    asyncio.run(main())