"""Speech resource for batch, streaming, and long-form TTS generation."""

from __future__ import annotations

import asyncio
import time
from collections.abc import AsyncIterator, Iterator
from contextlib import asynccontextmanager, contextmanager
from typing import TYPE_CHECKING, Any

from murmr._audio import concatenate_pcm_chunks, estimate_duration_ms
from murmr._chunker import split_into_chunks
from murmr._errors import MurmrChunkError, MurmrError
from murmr._streaming import (
    async_stream,
    collect_stream_as_pcm_async,
    collect_stream_as_pcm_sync,
    sync_stream,
)
from murmr._types import AsyncJobResponse, AudioStreamChunk, JobStatus, LongFormResult, SyncAudioResponse
from murmr._validate import validate_input

if TYPE_CHECKING:
    from collections.abc import Callable

    from murmr._client import AsyncMurmrClient, MurmrClient

DEFAULT_CHUNK_SIZE = 3500
DEFAULT_SILENCE_MS = 400
DEFAULT_MAX_RETRIES = 3
RATE_LIMIT_MIN_WAIT_S = 5.0


def _get_retry_delay_s(error: Exception, attempt: int) -> float:
    """Calculate retry delay with exponential backoff."""
    base: float = float(2**attempt)
    if isinstance(error, MurmrError) and error.status == 429:
        return max(RATE_LIMIT_MIN_WAIT_S, base)
    return base


class AsyncSpeech:
    """Async resource for speech generation."""

    def __init__(self, client: AsyncMurmrClient) -> None:
        self._client = client

    async def create(
        self,
        *,
        input: str,
        voice: str,
        voice_clone_prompt: str | None = None,
        language: str = "English",
        response_format: str = "wav",
        webhook_url: str | None = None,
    ) -> SyncAudioResponse | AsyncJobResponse:
        """Generate speech from text using a saved voice.

        Default (sync): Returns SyncAudioResponse with audio bytes (HTTP 200).
        With webhook_url: Returns AsyncJobResponse with a job ID (HTTP 202).
        """
        validate_input(input)
        if not voice or not voice.strip():
            raise MurmrError("voice ID is required")

        if webhook_url is not None:
            _validate_webhook_url(webhook_url)

        voice_fields: dict[str, str] = (
            {"voice_clone_prompt": voice_clone_prompt}
            if voice_clone_prompt
            else {"voice": voice}
        )

        body: dict[str, Any] = {
            "input": input,
            **voice_fields,
            "language": language,
            "response_format": response_format,
        }
        if webhook_url is not None:
            body["webhook_url"] = webhook_url

        response = await self._client.request("POST", "/v1/audio/speech", json=body)

        if response.status_code == 200:
            return SyncAudioResponse(
                audio=response.content,
                content_type=response.headers.get("content-type", "audio/wav"),
                duration_ms=int(response.headers.get("x-audio-duration-ms", "0")),
                total_time_ms=int(response.headers.get("x-total-time-ms", "0")),
            )

        data: dict[str, Any] = response.json()
        return AsyncJobResponse.model_validate(data)

    async def create_and_wait(
        self,
        *,
        input: str,
        voice: str,
        voice_clone_prompt: str | None = None,
        language: str = "English",
        response_format: str = "wav",
        poll_interval_s: float = 3.0,
        timeout_s: float = 900.0,
        on_poll: Callable[[JobStatus], None] | None = None,
    ) -> SyncAudioResponse | JobStatus:
        """Submit a batch job and wait for completion.

        If the server returns audio synchronously (default), returns
        SyncAudioResponse immediately. Otherwise polls until completion.
        """
        result = await self.create(
            input=input,
            voice=voice,
            voice_clone_prompt=voice_clone_prompt,
            language=language,
            response_format=response_format,
        )

        if isinstance(result, SyncAudioResponse):
            return result

        return await self._client.jobs.wait_for_completion(
            result.id,
            poll_interval_s=poll_interval_s,
            timeout_s=timeout_s,
            on_poll=on_poll,
        )

    @asynccontextmanager
    async def stream(
        self,
        *,
        input: str,
        voice: str,
        voice_clone_prompt: str | None = None,
        language: str = "English",
    ) -> AsyncIterator[AsyncIterator[AudioStreamChunk]]:
        """Stream speech from text using a saved voice.

        Usage::

            async with client.speech.stream(input="Hello", voice="voice_xxx") as stream:
                async for chunk in stream:
                    pcm = chunk.audio_bytes
        """
        validate_input(input)
        if not voice or not voice.strip():
            raise MurmrError("voice ID is required")

        voice_fields: dict[str, str] = (
            {"voice_clone_prompt": voice_clone_prompt}
            if voice_clone_prompt
            else {"voice": voice}
        )

        body = {"input": input, **voice_fields, "language": language}

        response = await self._client.request(
            "POST", "/v1/audio/speech/stream", json=body, stream=True
        )
        async with async_stream(response) as chunks:
            yield chunks

    async def create_long_form(
        self,
        *,
        input: str,
        voice: str,
        voice_clone_prompt: str | None = None,
        language: str = "English",
        chunk_size: int = DEFAULT_CHUNK_SIZE,
        silence_between_chunks_ms: int = DEFAULT_SILENCE_MS,
        max_retries: int = DEFAULT_MAX_RETRIES,
        start_from_chunk: int = 0,
        on_progress: Callable[[int, int, int], None] | None = None,
    ) -> LongFormResult:
        """Generate long-form audio from text of any length.

        Handles chunking, retries, progress reporting, and concatenation.

        Args:
            input: Text of any length.
            voice: Voice ID to use.
            voice_clone_prompt: Optional base64-encoded embedding data.
            language: Language for synthesis.
            chunk_size: Maximum characters per chunk.
            silence_between_chunks_ms: Silence to insert between chunks.
            max_retries: Max retry attempts per chunk.
            start_from_chunk: Resume from this chunk index.
            on_progress: Callback(current, total, percent) after each chunk.

        Returns:
            LongFormResult with concatenated WAV audio.
        """
        chunks = split_into_chunks(input, chunk_size)
        if not chunks:
            return LongFormResult(
                audio=b"", total_chunks=0, duration_ms=0, character_count=0
            )

        pcm_chunks: list[bytes] = []
        total_characters = 0

        for i, chunk_text in enumerate(chunks):
            if i < start_from_chunk:
                if on_progress is not None:
                    percent = round(((i + 1) / len(chunks)) * 100)
                    on_progress(i + 1, len(chunks), percent)
                continue

            last_error: Exception | None = None
            succeeded = False

            for retry in range(max_retries + 1):
                try:
                    pcm = await self._generate_chunk_audio(
                        chunk_text, voice, voice_clone_prompt, language
                    )
                    pcm_chunks.append(pcm)
                    total_characters += len(chunk_text)
                    succeeded = True

                    if on_progress is not None:
                        percent = round(((i + 1) / len(chunks)) * 100)
                        on_progress(i + 1, len(chunks), percent)
                    break
                except Exception as err:
                    last_error = err
                    if retry < max_retries:
                        await asyncio.sleep(_get_retry_delay_s(err, retry))

            if not succeeded:
                raise MurmrChunkError(
                    f"Chunk {i + 1}/{len(chunks)} failed after {max_retries + 1} attempts",
                    chunk_index=i,
                    completed_chunks=i,
                    total_chunks=len(chunks),
                    cause=last_error,
                )

        audio = concatenate_pcm_chunks(pcm_chunks, silence_between_chunks_ms)
        return LongFormResult(
            audio=audio,
            total_chunks=len(chunks),
            duration_ms=estimate_duration_ms(audio),
            character_count=total_characters,
        )

    async def _generate_chunk_audio(
        self,
        text: str,
        voice: str,
        voice_clone_prompt: str | None,
        language: str,
    ) -> bytes:
        """Generate PCM audio for a single chunk via the streaming endpoint."""
        voice_fields: dict[str, str] = (
            {"voice_clone_prompt": voice_clone_prompt}
            if voice_clone_prompt
            else {"voice": voice}
        )

        body = {"text": text, **voice_fields, "language": language}
        response = await self._client.request(
            "POST", "/v1/audio/speech/stream", json=body, stream=True
        )
        return await collect_stream_as_pcm_async(response)


class SyncSpeech:
    """Sync resource for speech generation."""

    def __init__(self, client: MurmrClient) -> None:
        self._client = client

    def create(
        self,
        *,
        input: str,
        voice: str,
        voice_clone_prompt: str | None = None,
        language: str = "English",
        response_format: str = "wav",
        webhook_url: str | None = None,
    ) -> SyncAudioResponse | AsyncJobResponse:
        """Generate speech from text using a saved voice.

        Default (sync): Returns SyncAudioResponse with audio bytes (HTTP 200).
        With webhook_url: Returns AsyncJobResponse with a job ID (HTTP 202).
        """
        validate_input(input)
        if not voice or not voice.strip():
            raise MurmrError("voice ID is required")

        if webhook_url is not None:
            _validate_webhook_url(webhook_url)

        voice_fields: dict[str, str] = (
            {"voice_clone_prompt": voice_clone_prompt}
            if voice_clone_prompt
            else {"voice": voice}
        )

        body: dict[str, Any] = {
            "input": input,
            **voice_fields,
            "language": language,
            "response_format": response_format,
        }
        if webhook_url is not None:
            body["webhook_url"] = webhook_url

        response = self._client.request("POST", "/v1/audio/speech", json=body)

        if response.status_code == 200:
            return SyncAudioResponse(
                audio=response.content,
                content_type=response.headers.get("content-type", "audio/wav"),
                duration_ms=int(response.headers.get("x-audio-duration-ms", "0")),
                total_time_ms=int(response.headers.get("x-total-time-ms", "0")),
            )

        data: dict[str, Any] = response.json()
        return AsyncJobResponse.model_validate(data)

    def create_and_wait(
        self,
        *,
        input: str,
        voice: str,
        voice_clone_prompt: str | None = None,
        language: str = "English",
        response_format: str = "wav",
        poll_interval_s: float = 3.0,
        timeout_s: float = 900.0,
        on_poll: Callable[[JobStatus], None] | None = None,
    ) -> SyncAudioResponse | JobStatus:
        """Submit a batch job and wait for completion.

        If the server returns audio synchronously (default), returns
        SyncAudioResponse immediately. Otherwise polls until completion.
        """
        result = self.create(
            input=input,
            voice=voice,
            voice_clone_prompt=voice_clone_prompt,
            language=language,
            response_format=response_format,
        )

        if isinstance(result, SyncAudioResponse):
            return result

        return self._client.jobs.wait_for_completion(
            result.id,
            poll_interval_s=poll_interval_s,
            timeout_s=timeout_s,
            on_poll=on_poll,
        )

    @contextmanager
    def stream(
        self,
        *,
        input: str,
        voice: str,
        voice_clone_prompt: str | None = None,
        language: str = "English",
    ) -> Iterator[Iterator[AudioStreamChunk]]:
        """Stream speech from text using a saved voice."""
        validate_input(input)
        if not voice or not voice.strip():
            raise MurmrError("voice ID is required")

        voice_fields: dict[str, str] = (
            {"voice_clone_prompt": voice_clone_prompt}
            if voice_clone_prompt
            else {"voice": voice}
        )

        body = {"input": input, **voice_fields, "language": language}

        response = self._client.request(
            "POST", "/v1/audio/speech/stream", json=body, stream=True
        )
        with sync_stream(response) as chunks:
            yield chunks

    def create_long_form(
        self,
        *,
        input: str,
        voice: str,
        voice_clone_prompt: str | None = None,
        language: str = "English",
        chunk_size: int = DEFAULT_CHUNK_SIZE,
        silence_between_chunks_ms: int = DEFAULT_SILENCE_MS,
        max_retries: int = DEFAULT_MAX_RETRIES,
        start_from_chunk: int = 0,
        on_progress: Callable[[int, int, int], None] | None = None,
    ) -> LongFormResult:
        """Generate long-form audio from text of any length (sync)."""
        chunks = split_into_chunks(input, chunk_size)
        if not chunks:
            return LongFormResult(
                audio=b"", total_chunks=0, duration_ms=0, character_count=0
            )

        pcm_chunks: list[bytes] = []
        total_characters = 0

        for i, chunk_text in enumerate(chunks):
            if i < start_from_chunk:
                if on_progress is not None:
                    percent = round(((i + 1) / len(chunks)) * 100)
                    on_progress(i + 1, len(chunks), percent)
                continue

            last_error: Exception | None = None
            succeeded = False

            for retry in range(max_retries + 1):
                try:
                    pcm = self._generate_chunk_audio(
                        chunk_text, voice, voice_clone_prompt, language
                    )
                    pcm_chunks.append(pcm)
                    total_characters += len(chunk_text)
                    succeeded = True

                    if on_progress is not None:
                        percent = round(((i + 1) / len(chunks)) * 100)
                        on_progress(i + 1, len(chunks), percent)
                    break
                except Exception as err:
                    last_error = err
                    if retry < max_retries:
                        time.sleep(_get_retry_delay_s(err, retry))

            if not succeeded:
                raise MurmrChunkError(
                    f"Chunk {i + 1}/{len(chunks)} failed after {max_retries + 1} attempts",
                    chunk_index=i,
                    completed_chunks=i,
                    total_chunks=len(chunks),
                    cause=last_error,
                )

        audio = concatenate_pcm_chunks(pcm_chunks, silence_between_chunks_ms)
        return LongFormResult(
            audio=audio,
            total_chunks=len(chunks),
            duration_ms=estimate_duration_ms(audio),
            character_count=total_characters,
        )

    def _generate_chunk_audio(
        self,
        text: str,
        voice: str,
        voice_clone_prompt: str | None,
        language: str,
    ) -> bytes:
        """Generate PCM audio for a single chunk via the streaming endpoint."""
        voice_fields: dict[str, str] = (
            {"voice_clone_prompt": voice_clone_prompt}
            if voice_clone_prompt
            else {"voice": voice}
        )

        body = {"text": text, **voice_fields, "language": language}
        response = self._client.request(
            "POST", "/v1/audio/speech/stream", json=body, stream=True
        )
        return collect_stream_as_pcm_sync(response)


def _validate_webhook_url(url: str) -> None:
    """Validate that a webhook URL is valid HTTPS."""
    from urllib.parse import urlparse

    try:
        parsed = urlparse(url)
        if parsed.scheme != "https":
            raise MurmrError("webhook_url must use HTTPS")
        if not parsed.netloc:
            raise MurmrError("webhook_url is not a valid URL")
    except MurmrError:
        raise
    except Exception:
        raise MurmrError("webhook_url is not a valid URL")
