from __future__ import annotations

import asyncio
from fractions import Fraction
from collections import deque

import numpy as np
from aiortc import MediaStreamTrack
from av import AudioFrame, VideoFrame

from core.state_machine import Arbitrator, SessionState
from services.avatar import AvatarService


class AudioBus:
    def __init__(self, sample_rate: int = 48000) -> None:
        self.sample_rate = sample_rate
        self._chunks: deque[np.ndarray] = deque()
        self._lock = asyncio.Lock()
        self._level = 0.0

    async def enqueue(self, pcm_int16_mono: np.ndarray) -> None:
        async with self._lock:
            if pcm_int16_mono.ndim != 1:
                pcm_int16_mono = pcm_int16_mono.reshape(-1)
            chunk = pcm_int16_mono.astype(np.int16, copy=False)
            self._chunks.append(chunk)
            if chunk.size > 0:
                rms = float(np.sqrt(np.mean((chunk.astype(np.float32) / 32768.0) ** 2)))
                self._level = 0.85 * self._level + 0.15 * min(1.0, rms * 10.0)

    async def read(self, samples: int) -> np.ndarray:
        out = np.zeros(samples, dtype=np.int16)
        async with self._lock:
            i = 0
            while i < samples and self._chunks:
                head = self._chunks[0]
                n = min(samples - i, head.shape[0])
                out[i : i + n] = head[:n]
                i += n
                if n == head.shape[0]:
                    self._chunks.popleft()
                else:
                    self._chunks[0] = head[n:]
            self._level = 0.92 * self._level
        return out

    async def level(self) -> float:
        async with self._lock:
            return float(self._level)

    async def clear(self) -> None:
        async with self._lock:
            self._chunks.clear()
            self._level = 0.0


class AvatarVideoTrack(MediaStreamTrack):
    kind = "video"

    def __init__(
        self,
        avatar_service: AvatarService,
        arbitrator: Arbitrator,
        audio_bus: AudioBus,
        fps: int = 25,
    ) -> None:
        super().__init__()
        self.avatar_service = avatar_service
        self.arbitrator = arbitrator
        self.audio_bus = audio_bus
        self.fps = fps
        self._pts = 0
        self._time_base = Fraction(1, 90000)

    async def recv(self) -> VideoFrame:
        await asyncio.sleep(1 / self.fps)
        generated = await self.avatar_service.pop_generated_frame()
        if generated is not None:
            arr = generated
        else:
            level = await self.audio_bus.level()
            speaking = self.arbitrator.state == SessionState.AVATAR_SPEAKING or level > 0.02
            arr = await self.avatar_service.render_frame(mouth_open=level, speaking=speaking)

        frame = VideoFrame.from_ndarray(arr, format="bgr24")
        frame.pts = self._pts
        frame.time_base = self._time_base
        self._pts += int(90000 / self.fps)
        return frame


class AvatarAudioTrack(MediaStreamTrack):
    kind = "audio"

    def __init__(self, audio_bus: AudioBus, sample_rate: int = 48000, frame_ms: int = 20) -> None:
        super().__init__()
        self.audio_bus = audio_bus
        self.sample_rate = sample_rate
        self.samples_per_frame = int(sample_rate * frame_ms / 1000)
        self._pts = 0
        self._time_base = Fraction(1, sample_rate)

    async def recv(self) -> AudioFrame:
        await asyncio.sleep(self.samples_per_frame / self.sample_rate)
        mono = await self.audio_bus.read(self.samples_per_frame)
        samples = mono.reshape(1, -1)
        frame = AudioFrame(format="s16", layout="mono", samples=self.samples_per_frame)
        frame.planes[0].update(samples.tobytes())
        frame.sample_rate = self.sample_rate
        frame.pts = self._pts
        frame.time_base = self._time_base
        self._pts += self.samples_per_frame
        return frame