from __future__ import annotations

import asyncio
import os
from typing import AsyncIterator, Optional

import azure.cognitiveservices.speech as speechsdk

from app.observability.logging import get_logger
from app.providers.tts_base import TTSProvider


def _wrap_with_sales_ssml(text: str, voice: str = "en-US-AriaNeural") -> str:
    """Wrap text with SSML for optimized sales voice styling."""
    # Don't double-wrap if already SSML
    if text.strip().startswith("<speak"):
        return text
    
    # Use customerservice style for AriaNeural, friendly for others
    style = "customerservice" if "Aria" in voice else "friendly"
    style_degree = "1.2" if "Aria" in voice else "1.0"
    
    return f'''<speak version="1.0" xmlns="http://www.w3.org/2001/10/synthesis" xml:lang="en-US">
    <voice name="{voice}" style="{style}" styledegree="{style_degree}">
        {text}
    </voice>
</speak>'''


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)

    voice = os.getenv("SPEECH_VOICE", "en-US-AriaNeural")
    config.speech_synthesis_voice_name = voice
    config.set_speech_synthesis_output_format(
        speechsdk.SpeechSynthesisOutputFormat.Raw16Khz16BitMonoPcm
    )
    return config


class FoundryTTSProvider(TTSProvider):
    async def synthesize_stream(self, text: str, voice: Optional[str] = None) -> AsyncIterator[bytes]:
        logger = get_logger("tts")
        speech_config = _speech_config()
        if voice:
            speech_config.speech_synthesis_voice_name = voice

        queue: asyncio.Queue[Optional[bytes]] = asyncio.Queue()
        loop = asyncio.get_running_loop()
        callback = _StreamingOutputStream(loop, queue)
        stream = speechsdk.audio.PushAudioOutputStream(callback)
        audio_config = speechsdk.audio.AudioOutputConfig(stream=stream)
        synthesizer = speechsdk.SpeechSynthesizer(speech_config=speech_config, audio_config=audio_config)

        # Wrap text with sales-optimized SSML
        voice_name = voice or speech_config.speech_synthesis_voice_name
        ssml_text = _wrap_with_sales_ssml(text, voice_name)

        async def _run() -> None:
            if ssml_text.lstrip().startswith("<speak"):
                result = await asyncio.to_thread(synthesizer.speak_ssml_async(ssml_text).get)
            else:
                result = await asyncio.to_thread(synthesizer.speak_text_async(ssml_text).get)
            if result.reason != speechsdk.ResultReason.SynthesizingAudioCompleted:
                if result.reason == speechsdk.ResultReason.Canceled:
                    details = speechsdk.SpeechSynthesisCancellationDetails.from_result(result)
                    logger.warning(
                        "tts_canceled",
                        reason=str(details.reason),
                        error_code=str(details.error_code),
                        error_details=details.error_details,
                    )
                else:
                    logger.warning("tts_failed", reason=str(result.reason))
                callback.end()
                return
            callback.end()

        task = asyncio.create_task(_run())

        sent = 0
        while True:
            chunk = await queue.get()
            if chunk is None:
                break
            sent += 1
            yield chunk

        await task
        if sent == 0:
            logger.warning("tts_empty_audio")


class _StreamingOutputStream(speechsdk.audio.PushAudioOutputStreamCallback):
    def __init__(self, loop: asyncio.AbstractEventLoop, queue: asyncio.Queue[Optional[bytes]]) -> None:
        self._loop = loop
        self._queue = queue
        self._closed = False

    def write(self, audio_buffer: memoryview) -> int:
        data = bytes(audio_buffer)
        if data:
            self._loop.call_soon_threadsafe(self._queue.put_nowait, data)
        return len(audio_buffer)

    def close(self) -> None:
        self.end()

    def end(self) -> None:
        if self._closed:
            return
        self._closed = True
        self._loop.call_soon_threadsafe(self._queue.put_nowait, None)
