from __future__ import annotations

import asyncio
import math
import re
import struct
from abc import ABC, abstractmethod
from collections.abc import AsyncGenerator


class BaseSTT(ABC):
    @abstractmethod
    async def transcribe(self, audio_bytes: bytes) -> str:
        raise NotImplementedError


class BaseLLM(ABC):
    @abstractmethod
    async def generate_stream(self, text: str, context: list) -> AsyncGenerator[str, None]:
        raise NotImplementedError


class BaseTTS(ABC):
    @abstractmethod
    async def synthesize_stream(self, text: str) -> AsyncGenerator[bytes, None]:
        raise NotImplementedError


class MockSTT(BaseSTT):
    def __init__(
        self,
        *,
        latency_ms: int = 40,
        sample_rate_hz: int = 8000,
        scripted_transcripts: list[str] | None = None,
    ) -> None:
        self._latency_ms = max(latency_ms, 0)
        self._sample_rate_hz = sample_rate_hz
        self._scripted_transcripts = list(scripted_transcripts or [])
        self._call_count = 0

    async def transcribe(self, audio_bytes: bytes) -> str:
        await asyncio.sleep(self._latency_ms / 1000.0)
        self._call_count += 1
        if self._scripted_transcripts:
            return self._scripted_transcripts.pop(0)
        duration_ms = int(((len(audio_bytes) // 2) / float(self._sample_rate_hz)) * 1000.0)
        return f"mock user utterance {self._call_count} ({duration_ms} ms)"


class MockLLM(BaseLLM):
    def __init__(
        self,
        *,
        token_delay_ms: int = 35,
        scripted_responses: list[str] | None = None,
    ) -> None:
        self._token_delay_ms = max(token_delay_ms, 0)
        self._scripted_responses = list(scripted_responses or [])

    async def generate_stream(self, text: str, context: list) -> AsyncGenerator[str, None]:
        del context
        response_text = (
            self._scripted_responses.pop(0)
            if self._scripted_responses
            else f"Mock reply: {text}"
        )
        chunks = re.findall(r"\S+\s*", response_text) or [response_text]
        for chunk in chunks:
            await asyncio.sleep(self._token_delay_ms / 1000.0)
            yield chunk


class MockTTS(BaseTTS):
    def __init__(
        self,
        *,
        sample_rate_hz: int = 8000,
        chunk_duration_ms: int = 40,
        chunk_delay_ms: int = 15,
        tone_hz: float = 440.0,
        amplitude: int = 9000,
        milliseconds_per_word: int = 140,
    ) -> None:
        self._sample_rate_hz = sample_rate_hz
        self._chunk_duration_ms = max(chunk_duration_ms, 20)
        self._chunk_delay_ms = max(chunk_delay_ms, 0)
        self._tone_hz = tone_hz
        self._amplitude = amplitude
        self._milliseconds_per_word = max(milliseconds_per_word, 40)

    async def synthesize_stream(self, text: str) -> AsyncGenerator[bytes, None]:
        word_count = max(len(text.split()), 1)
        total_ms = min(word_count * self._milliseconds_per_word, 2400)
        total_samples = max(int(self._sample_rate_hz * (total_ms / 1000.0)), 1)
        chunk_samples = max(int(self._sample_rate_hz * (self._chunk_duration_ms / 1000.0)), 1)

        for start in range(0, total_samples, chunk_samples):
            end = min(start + chunk_samples, total_samples)
            pcm = bytearray()
            for index in range(start, end):
                angle = 2.0 * math.pi * self._tone_hz * (index / float(self._sample_rate_hz))
                sample = int(self._amplitude * math.sin(angle))
                pcm.extend(struct.pack("<h", sample))
            await asyncio.sleep(self._chunk_delay_ms / 1000.0)
            yield bytes(pcm)