from __future__ import annotations

import asyncio
import os
from typing import AsyncIterator, Optional

import azure.cognitiveservices.speech as speechsdk

from app.providers.stt_base import STTProvider, STTStream, TranscriptEvent


def _speech_config() -> speechsdk.SpeechConfig:
    key = os.getenv("SPEECH_KEY")
    endpoint = os.getenv("ENDPOINT") or os.getenv("SPEECH_ENDPOINT")
    region = os.getenv("SPEECH_REGION")

    if not key or not (endpoint or region):
        raise RuntimeError("Missing SPEECH_KEY and ENDPOINT or SPEECH_REGION for Foundry Speech.")

    if endpoint:
        config = speechsdk.SpeechConfig(subscription=key, endpoint=endpoint)
    else:
        config = speechsdk.SpeechConfig(subscription=key, region=region)

    language = os.getenv("SPEECH_LANGUAGE", "en-US")
    config.speech_recognition_language = language

    end_silence = os.getenv("SPEECH_END_SILENCE_MS", "600")
    initial_silence = os.getenv("SPEECH_INITIAL_SILENCE_MS", "4000")
    stable_partial = os.getenv("SPEECH_STABLE_PARTIAL", "3")
    segment_silence = os.getenv("SPEECH_SEGMENT_SILENCE_MS")

    try:
        config.set_property(speechsdk.PropertyId.SpeechServiceConnection_EndSilenceTimeoutMs, end_silence)
        config.set_property(speechsdk.PropertyId.SpeechServiceConnection_InitialSilenceTimeoutMs, initial_silence)
        config.set_property(speechsdk.PropertyId.SpeechServiceResponse_StablePartialResultThreshold, stable_partial)
        if segment_silence:
            config.set_property(
                speechsdk.PropertyId.SpeechServiceConnection_SegmentationSilenceTimeoutMs, segment_silence
            )
    except Exception:
        # If a property isn't supported in this SDK version, ignore.
        pass
    return config


class FoundrySTTStream(STTStream):
    def __init__(self, speech_config: speechsdk.SpeechConfig, sample_rate: int) -> None:
        self._loop = asyncio.get_running_loop()
        self._queue: asyncio.Queue[Optional[TranscriptEvent]] = asyncio.Queue()
        self._done = asyncio.Event()

        stream_format = speechsdk.audio.AudioStreamFormat(samples_per_second=sample_rate, bits_per_sample=16, channels=1)
        self._push_stream = speechsdk.audio.PushAudioInputStream(stream_format=stream_format)
        audio_config = speechsdk.audio.AudioConfig(stream=self._push_stream)
        self._recognizer = speechsdk.SpeechRecognizer(speech_config=speech_config, audio_config=audio_config)

        self._recognizer.recognizing.connect(self._on_recognizing)
        self._recognizer.recognized.connect(self._on_recognized)
        self._recognizer.canceled.connect(self._on_canceled)
        self._recognizer.session_stopped.connect(self._on_session_stopped)

        self._recognizer.start_continuous_recognition_async()

    def _emit(self, event: Optional[TranscriptEvent]) -> None:
        self._loop.call_soon_threadsafe(self._queue.put_nowait, event)

    def _on_recognizing(self, evt: speechsdk.SpeechRecognitionEventArgs) -> None:
        text = evt.result.text
        if text:
            self._emit(TranscriptEvent(text=text, is_final=False))

    def _on_recognized(self, evt: speechsdk.SpeechRecognitionEventArgs) -> None:
        if evt.result.reason == speechsdk.ResultReason.RecognizedSpeech:
            text = evt.result.text
            if text:
                self._emit(TranscriptEvent(text=text, is_final=True))

    def _on_canceled(self, _: speechsdk.SpeechRecognitionEventArgs) -> None:
        self._finish()

    def _on_session_stopped(self, _: speechsdk.SessionEventArgs) -> None:
        self._finish()

    def _finish(self) -> None:
        if self._done.is_set():
            return
        self._done.set()
        self._emit(None)

    async def send_audio(self, data: bytes) -> None:
        if data:
            self._push_stream.write(data)

    async def end(self) -> None:
        self._push_stream.close()
        await asyncio.to_thread(self._recognizer.stop_continuous_recognition_async().get)
        self._finish()

    async def results(self) -> AsyncIterator[TranscriptEvent]:
        while True:
            item = await self._queue.get()
            if item is None:
                break
            yield item


class FoundrySTTProvider(STTProvider):
    async def start_stream(self, sample_rate: int) -> STTStream:
        speech_config = _speech_config()
        return FoundrySTTStream(speech_config, sample_rate)
