"""Live animated face for the mask, driven by pre-uploaded DIY frames + PLAY.

This is the smooth path (no per-frame upload logo): a frame set is uploaded to
the mask once (slow, one-time), then animation is just ``PLAY <slot>`` commands
(fast). The background loop blinks, glances around when idle, and moves the
mouth while "speaking".

Driving the mouth (best -> simplest):
  * :meth:`set_audio_level` (0..1 from live audio RMS) -> real lip-sync; the
    mouth opens proportionally to loudness and closes on pauses. Use this with
    Gemini Live / TTS audio for a human-looking mouth.
  * :meth:`set_speaking` (True/False) -> a natural *auto* talk envelope (varied
    open/close with word-break pauses) when you don't have amplitude.
  * :meth:`set_mouth` (0..3) -> drive a fixed level directly.

Interactivity:
  * :meth:`set_expression` holds an expression (e.g. "surprised", "smile",
    "listening") over the idle/talk animation until cleared with ``None``.
"""

from __future__ import annotations

import asyncio
import random
import time
from typing import Dict, Optional

import colorface

_GLANCE = ["look_left", "neutral", "look_right", "neutral"]
_MOUTH_BY_LEVEL = ["neutral", "talk1", "talk2", "talk3"]


class FaceAnimator:
    def __init__(self, mask, frames: Optional[Dict[str, bytes]] = None, *,
                 fps: float = 8.0, brightness: int = 95):
        self.mask = mask
        self.frames = frames or colorface.default_frames()
        self.fps = fps
        self.brightness = brightness
        self.slots: Dict[str, int] = {}

        # drive state
        self._speaking = False          # auto talk envelope
        self._mouth: Optional[int] = None   # fixed manual level
        self._level = 0.0               # live audio level 0..1 (lip-sync)
        self._lipsync = False
        self._expression: Optional[str] = None  # held override frame
        self._blink = True

        # internal animation state
        self._talk_cur = 0.0            # smoothed auto-talk mouth position
        self._talk_target = 0.0
        self._talk_hold = 0

        self._task: Optional[asyncio.Task] = None
        self._stop = False
        self._last: Optional[str] = None

    # -- lifecycle ------------------------------------------------------------

    async def load(self, force: bool = False):
        """Upload the frame set to the mask (skipped if already present).

        Resilient to a flaky BLE link: if the connection drops mid-upload, it
        reconnects and retries the current frame, so a weak signal slows the
        one-time upload rather than failing it.
        """
        names = list(self.frames)
        count = await self.mask.get_diy_count(timeout=4.0) or 0
        if not force and count >= len(names):
            # Frames already stored (DIY images persist on the mask's flash).
            # ``>=`` tolerates the CHEC count being off-by-one on some units, so
            # we don't needlessly re-upload (which is slow + risks corruption).
            self.slots = {name: i + 1 for i, name in enumerate(names)}
            return
        await self.mask.clear_diy()
        for i, (name, data) in enumerate(self.frames.items(), start=1):
            await self._upload_frame_resilient(i, data)
            self.slots[name] = i
            await asyncio.sleep(0.2)

    async def _upload_frame_resilient(self, slot: int, data: bytes, tries: int = 5):
        from bleak.exc import BleakError
        for attempt in range(tries):
            try:
                await self.mask.upload_frame(data, slot)
                return
            except (BleakError, EOFError, OSError) as exc:
                if attempt == tries - 1:
                    raise
                # link dropped mid-upload -> reconnect and re-send this frame
                # (a fresh DATS resets any half-finished slot on the mask)
                try:
                    await self.mask.disconnect()
                except Exception:
                    pass
                await asyncio.sleep(1.0)
                try:
                    await self.mask.connect(timeout=15.0, attempts=8)
                    await self.mask.set_brightness(self.brightness)
                except Exception:
                    pass  # next attempt's upload_frame will surface a clear error

    async def start(self, *, reload: bool = False):
        await self.mask.set_brightness(self.brightness)
        await self.load(force=reload)
        self._stop = False
        await self._play("neutral")
        self._task = asyncio.create_task(self._loop())
        return self

    async def stop(self):
        self._stop = True
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None

    async def __aenter__(self):
        return await self.start()

    async def __aexit__(self, *exc):
        await self.stop()

    # -- control --------------------------------------------------------------

    def set_audio_level(self, level: float):
        """Real lip-sync: 0..1 live audio amplitude -> mouth openness.

        Call this as fast as you like from an audio callback; the loop samples
        it at ``fps`` and the level decays so the mouth closes on silence.
        """
        self._lipsync = True
        self._speaking = False
        self._mouth = None
        self._level = max(self._level * 0.4, min(1.0, float(level)))  # fast attack

    def set_speaking(self, speaking: bool):
        """Auto talk envelope while speaking; closes the mouth when False."""
        self._mouth = None
        self._lipsync = False
        self._speaking = bool(speaking)
        if not speaking:
            self._talk_cur = self._talk_target = 0.0

    def set_mouth(self, level: int):
        """Drive the mouth at a fixed level (0=closed..3=wide)."""
        self._speaking = False
        self._lipsync = False
        self._mouth = max(0, min(3, int(level)))

    def set_expression(self, name: Optional[str]):
        """Hold an expression frame over the animation (None clears it)."""
        self._expression = name if name in self.frames else None

    def set_blink(self, enabled: bool):
        self._blink = bool(enabled)

    async def show(self, name: str):
        """Show one named frame immediately (one-off)."""
        await self._play(name)

    # -- internals ------------------------------------------------------------

    @staticmethod
    def _level_to_mouth(level: float) -> int:
        if level < 0.06:
            return 0
        if level < 0.16:
            return 1
        if level < 0.32:
            return 2
        return 3

    def _auto_talk_mouth(self) -> int:
        """A speech-like envelope: ease toward random targets, with word breaks."""
        if self._talk_hold <= 0:
            # mostly open of varying degree, with occasional full closes (breaks)
            self._talk_target = random.choice([0, 1, 1, 2, 2, 3, 3, 2, 1, 0])
            self._talk_hold = random.randint(1, 3)
        self._talk_hold -= 1
        # ease current toward target so motion is smooth, not steppy
        self._talk_cur += (self._talk_target - self._talk_cur) * 0.6
        return int(round(self._talk_cur))

    async def _play(self, name: str):
        slot = self.slots.get(name)
        if slot is None or name == self._last:
            return
        try:
            await self.mask.play_diy(slot)
            self._last = name
        except Exception:
            self._last = None  # retry next tick on transient BLE error

    async def _loop(self):
        dt = 1.0 / self.fps
        next_blink = time.monotonic() + random.uniform(2.0, 5.0)
        next_glance = time.monotonic() + random.uniform(4.0, 9.0)
        while not self._stop:
            now = time.monotonic()

            # periodic blink (skipped while a held expression is active)
            if self._blink and self._expression is None and now >= next_blink:
                await self._play("blink")
                await asyncio.sleep(0.12)
                next_blink = now + random.uniform(2.5, 6.0)

            if self._expression is not None:
                await self._play(self._expression)
            elif self._lipsync:
                await self._play(_MOUTH_BY_LEVEL[self._level_to_mouth(self._level)])
                self._level *= 0.55                       # release toward closed
                if self._level < 0.02:
                    self._lipsync = False                 # idle once silent
            elif self._speaking:
                await self._play(_MOUTH_BY_LEVEL[self._auto_talk_mouth()])
            elif self._mouth is not None:
                await self._play(_MOUTH_BY_LEVEL[self._mouth])
            elif now >= next_glance:
                for g in _GLANCE:
                    await self._play(g)
                    await asyncio.sleep(0.25)
                next_glance = now + random.uniform(5.0, 11.0)
            else:
                await self._play("neutral")

            await asyncio.sleep(dt)