from __future__ import annotations

import asyncio
import math
import time
from typing import Any

import numpy as np

from config import settings


class AvatarService:
    def __init__(self) -> None:
        self.driver_name = "arkit-blendshape-v1"
        self._lock = asyncio.Lock()
        self._sequence = 0
        self._last_error: str | None = None
        self._last_frame_count = 0
        self._last_chunk_duration_ms = 0

    async def build_idle_payload(self, reason: str = "idle") -> dict[str, Any]:
        frame = self._build_frame(seq=await self._next_sequence(), time_ms=0, controls=self._neutral_controls())
        self._last_frame_count = 1
        self._last_chunk_duration_ms = 0
        self._last_error = None
        return {
            "type": "animation_state",
            "driver": self.driver_name,
            "mode": settings.avatar_driver_mode,
            "schema": settings.avatar_blendshape_schema,
            "fps": settings.avatar_fps,
            "reason": reason,
            "frame_count": 1,
            "frames": [frame],
            "ts_ms": int(time.time() * 1000),
        }

    async def build_reset_payload(self, reason: str = "reset") -> dict[str, Any]:
        payload = await self.build_idle_payload(reason=reason)
        payload["type"] = "animation_reset"
        return payload

    async def build_controls_from_audio(
        self,
        audio: np.ndarray,
        sr: int,
        *,
        text: str = "",
        chunk_index: int = 0,
        total_chunks: int = 1,
    ) -> dict[str, Any]:
        if sr <= 0:
            self._last_error = f"invalid sample rate: {sr}"
            return await self.build_idle_payload(reason="invalid-sample-rate")

        mono = np.asarray(audio, dtype=np.float32).reshape(-1)
        if mono.size <= 1:
            return await self.build_idle_payload(reason="empty-audio")

        mono = np.clip(mono, -1.0, 1.0)
        frame_count = max(1, int(math.ceil((mono.shape[0] / float(sr)) * settings.avatar_fps)))
        seqs = await self._reserve_sequences(frame_count)
        frames = self._audio_to_frames(mono, sr, seqs)

        self._last_frame_count = len(frames)
        self._last_chunk_duration_ms = int(round(1000.0 * mono.shape[0] / float(sr)))
        self._last_error = None
        return {
            "type": "animation_chunk",
            "driver": self.driver_name,
            "mode": settings.avatar_driver_mode,
            "schema": settings.avatar_blendshape_schema,
            "chunk_index": chunk_index,
            "total_chunks": total_chunks,
            "sample_rate": sr,
            "fps": settings.avatar_fps,
            "text": text,
            "frame_count": len(frames),
            "duration_ms": self._last_chunk_duration_ms,
            "frames": frames,
            "ts_ms": int(time.time() * 1000),
        }

    @property
    def schema(self) -> dict[str, Any]:
        return {
            "driver": self.driver_name,
            "mode": settings.avatar_driver_mode,
            "protocol": settings.avatar_control_protocol,
            "schema": settings.avatar_blendshape_schema,
            "controls": [
                "jawOpen",
                "mouthClose",
                "mouthFunnel",
                "mouthPucker",
                "viseme_aa",
                "viseme_ee",
                "viseme_oh",
                "headYaw",
                "headPitch",
                "headRoll",
            ],
        }

    def _audio_to_frames(self, audio: np.ndarray, sr: int, sequences: range) -> list[dict[str, Any]]:
        samples_per_frame = max(1, int(round(sr / max(1, settings.avatar_fps))))
        frames: list[dict[str, Any]] = []
        jaw_prev = 0.0
        for index, seq in enumerate(sequences):
            start = index * samples_per_frame
            end = min(audio.shape[0], start + samples_per_frame)
            window = audio[start:end]
            if window.size == 0:
                window = np.zeros(1, dtype=np.float32)
            energy, brightness = self._extract_features(window, sr)
            jaw_open = 0.70 * jaw_prev + 0.30 * energy
            jaw_prev = jaw_open
            phase = index / max(1, len(sequences) - 1)
            frames.append(
                self._build_frame(
                    seq=seq,
                    time_ms=int(round(index * 1000.0 / settings.avatar_fps)),
                    controls=self._controls_from_features(jaw_open, brightness, phase),
                )
            )
        return frames

    def _extract_features(self, window: np.ndarray, sr: int) -> tuple[float, float]:
        rms = float(np.sqrt(np.mean(window**2))) if window.size else 0.0
        energy = float(np.clip((rms - 0.01) * 10.0, 0.0, 1.0))
        if window.size < 16:
            return energy, 0.5
        spectrum = np.abs(np.fft.rfft(window * np.hanning(window.size)))
        if spectrum.size <= 1:
            return energy, 0.5
        freqs = np.fft.rfftfreq(window.size, d=1.0 / float(sr))
        weights = spectrum[1:]
        freq_values = freqs[1:]
        denom = float(np.sum(weights)) + 1e-6
        centroid = float(np.sum(freq_values * weights) / denom)
        brightness = float(np.clip((centroid - 500.0) / 2500.0, 0.0, 1.0))
        return energy, brightness

    def _controls_from_features(self, energy: float, brightness: float, phase: float) -> dict[str, float]:
        jaw_open = float(np.clip(energy, 0.0, 1.0))
        mouth_close = float(np.clip(1.0 - jaw_open * 0.82, 0.0, 1.0))
        mouth_funnel = float(np.clip(0.16 + jaw_open * (0.40 - brightness * 0.12), 0.0, 1.0))
        mouth_pucker = float(np.clip(0.06 + jaw_open * (0.20 + (1.0 - brightness) * 0.30), 0.0, 1.0))
        viseme_aa = jaw_open
        viseme_ee = float(np.clip(jaw_open * (0.25 + brightness * 0.95), 0.0, 1.0))
        viseme_oh = float(np.clip(jaw_open * (0.30 + (1.0 - brightness) * 0.90), 0.0, 1.0))
        head_yaw = float(np.clip(math.sin(phase * math.pi * 1.7) * 0.02 * max(jaw_open, 0.18), -0.25, 0.25))
        head_pitch = float(np.clip(math.cos(phase * math.pi * 2.0) * 0.03 * max(jaw_open, 0.15), -0.25, 0.25))
        head_roll = float(np.clip(math.sin(phase * math.pi) * 0.015, -0.15, 0.15))
        return {
            "jawOpen": round(jaw_open, 4),
            "mouthClose": round(mouth_close, 4),
            "mouthFunnel": round(mouth_funnel, 4),
            "mouthPucker": round(mouth_pucker, 4),
            "viseme_aa": round(viseme_aa, 4),
            "viseme_ee": round(viseme_ee, 4),
            "viseme_oh": round(viseme_oh, 4),
            "headYaw": round(head_yaw, 4),
            "headPitch": round(head_pitch, 4),
            "headRoll": round(head_roll, 4),
        }

    def _build_frame(self, seq: int, time_ms: int, controls: dict[str, float]) -> dict[str, Any]:
        return {
            "seq": seq,
            "time_ms": time_ms,
            "speaking": controls["jawOpen"] > 0.04,
            "controls": controls,
        }

    def _neutral_controls(self) -> dict[str, float]:
        return {
            "jawOpen": 0.0,
            "mouthClose": 1.0,
            "mouthFunnel": 0.0,
            "mouthPucker": 0.0,
            "viseme_aa": 0.0,
            "viseme_ee": 0.0,
            "viseme_oh": 0.0,
            "headYaw": 0.0,
            "headPitch": 0.0,
            "headRoll": 0.0,
        }

    async def _next_sequence(self) -> int:
        async with self._lock:
            seq = self._sequence
            self._sequence += 1
        return seq

    async def _reserve_sequences(self, count: int) -> range:
        async with self._lock:
            start = self._sequence
            self._sequence += count
        return range(start, start + count)

    @property
    def health(self) -> dict[str, Any]:
        return {
            "driver": self.driver_name,
            "mode": settings.avatar_driver_mode,
            "protocol": settings.avatar_control_protocol,
            "schema": settings.avatar_blendshape_schema,
            "last_frame_count": self._last_frame_count,
            "last_chunk_duration_ms": self._last_chunk_duration_ms,
            "last_error": self._last_error,
        }