from __future__ import annotations

from pathlib import Path
from typing import Any

from .client import (
    DEFAULT_BASE_URL,
    DEFAULT_CHANNELS,
    DEFAULT_CHUNK_SIZE,
    DEFAULT_MODEL,
    DEFAULT_SAMPLE_RATE,
    DEFAULT_SAMPLE_WIDTH,
    DEFAULT_TIMEOUT,
    SDK_VERSION,
    SpeechResult,
)
from .exceptions import OuteAIAPIError, OuteAIError


class AsyncSpeechStream:
    def __init__(self, stream_context: Any, *, error_handler: Any, chunk_size: int = DEFAULT_CHUNK_SIZE) -> None:
        self._stream_context = stream_context
        self._error_handler = error_handler
        self._chunk_size = chunk_size
        self._response: Any | None = None
        self._closed = False

    @property
    def headers(self) -> Any:
        if self._response is None:
            raise OuteAIError("The async stream has not been opened yet.")
        return self._response.headers

    @property
    def status_code(self) -> int:
        if self._response is None:
            raise OuteAIError("The async stream has not been opened yet.")
        return self._response.status_code

    async def _open(self) -> None:
        if self._response is None:
            self._response = await self._stream_context.__aenter__()
            await self._error_handler(self._response)

    async def close(self) -> None:
        if not self._closed:
            self._closed = True
            await self._stream_context.__aexit__(None, None, None)

    async def __aenter__(self) -> "AsyncSpeechStream":
        await self._open()
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.close()

    def __aiter__(self):
        return self._iterate()

    async def _iterate(self):
        await self._open()
        try:
            async for chunk in self._response.aiter_bytes(self._chunk_size):
                if chunk:
                    yield chunk
        finally:
            await self.close()

    async def save(self, path: str | Path) -> Path:
        output_path = Path(path)
        with output_path.open("wb") as handle:
            async for chunk in self:
                handle.write(chunk)
        return output_path


class AsyncOuteAI:
    def __init__(
        self,
        api_key: str,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timeout: int | float = DEFAULT_TIMEOUT,
        session: Any | None = None,
    ) -> None:
        api_key = api_key.strip() if isinstance(api_key, str) else ""
        if not api_key:
            raise OuteAIError("`api_key` must be a non-empty string.")

        try:
            import httpx  # type: ignore
        except ImportError as exc:
            raise OuteAIError("Async support requires `pip install 'outeai[async]'`.") from exc

        self._httpx = httpx
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self.session = session or httpx.AsyncClient(
            headers={
                "Authorization": f"Bearer {api_key}",
                "Accept": "application/json",
                "User-Agent": f"outeai-python-sdk/{SDK_VERSION}",
            },
            timeout=timeout,
        )

    async def __aenter__(self) -> "AsyncOuteAI":
        return self

    async def __aexit__(self, exc_type, exc, tb) -> None:
        await self.close()

    async def close(self) -> None:
        await self.session.aclose()

    async def health(self) -> dict[str, Any]:
        response = await self._request("GET", "/api/v1/health")
        return await self._json(response)

    async def list_voices(self) -> list[dict[str, Any]]:
        response = await self._request("GET", "/api/v1/voices")
        data = await self._data(response)
        return data.get("voices", [])

    async def generate_speech(
        self,
        *,
        text: str,
        voice_id: str,
        model: str = DEFAULT_MODEL,
    ) -> SpeechResult:
        payload = self._tts_payload(text=text, voice_id=voice_id, model=model)
        response = await self._request(
            "POST",
            "/api/v1/tts_non_stream",
            json=payload,
            headers={"Accept": "audio/wav"},
        )
        try:
            wav_bytes = await response.aread()
            content_type = response.headers.get("Content-Type", "audio/wav")
        finally:
            await response.aclose()

        return SpeechResult(
            wav_bytes=wav_bytes,
            content_type=content_type,
            sample_rate=DEFAULT_SAMPLE_RATE,
            channels=DEFAULT_CHANNELS,
            sample_width=DEFAULT_SAMPLE_WIDTH,
        )

    async def generate_speech_to_file(
        self,
        path: str | Path,
        *,
        text: str,
        voice_id: str,
        model: str = DEFAULT_MODEL,
    ) -> Path:
        result = await self.generate_speech(text=text, voice_id=voice_id, model=model)
        return result.save(path)

    def stream_speech(
        self,
        *,
        text: str,
        voice_id: str,
        model: str = DEFAULT_MODEL,
        chunk_size: int = DEFAULT_CHUNK_SIZE,
    ) -> AsyncSpeechStream:
        payload = self._tts_payload(text=text, voice_id=voice_id, model=model)
        stream_context = self.session.stream(
            "POST",
            f"{self.base_url}/api/v1/tts_stream",
            json=payload,
            headers={"Accept": "audio/wav"},
        )
        return AsyncSpeechStream(
            stream_context,
            error_handler=self._raise_for_error,
            chunk_size=chunk_size,
        )

    async def stream_speech_to_file(
        self,
        path: str | Path,
        *,
        text: str,
        voice_id: str,
        model: str = DEFAULT_MODEL,
        chunk_size: int = DEFAULT_CHUNK_SIZE,
    ) -> Path:
        async with self.stream_speech(
            text=text,
            voice_id=voice_id,
            model=model,
            chunk_size=chunk_size,
        ) as stream:
            return await stream.save(path)

    def _tts_payload(self, *, text: str, voice_id: str, model: str) -> dict[str, str]:
        text = text.strip() if isinstance(text, str) else ""
        voice_id = voice_id.strip() if isinstance(voice_id, str) else ""
        model = model.strip() if isinstance(model, str) else ""

        if not text:
            raise OuteAIError("`text` must be a non-empty string.")
        if not voice_id:
            raise OuteAIError("`voice_id` must be a non-empty string.")
        if not model:
            raise OuteAIError("`model` must be a non-empty string.")

        return {
            "text": text,
            "voice_id": voice_id,
            "model_name": model,
            "response_format": "wav",
        }

    async def _request(
        self,
        method: str,
        path: str,
        *,
        headers: dict[str, str] | None = None,
        **kwargs: Any,
    ) -> Any:
        try:
            response = await self.session.request(
                method=method,
                url=f"{self.base_url}{path}",
                headers=dict(headers or {}),
                **kwargs,
            )
        except self._httpx.HTTPError as exc:
            raise OuteAIError(f"Request to {self.base_url}{path} failed: {exc}") from exc

        await self._raise_for_error(response)
        return response

    async def _raise_for_error(self, response: Any) -> None:
        if response.is_success:
            return

        payload: dict[str, Any] | None = None
        code = None

        try:
            payload = response.json()
        except ValueError:
            text = await response.aread()
            message = text.decode("utf-8", errors="replace").strip()
        else:
            if isinstance(payload, dict):
                code = payload.get("code")
                message = str(payload.get("message") or "")
            else:
                message = ""

        if not message:
            message = f"Request failed with status {response.status_code}."

        raise OuteAIAPIError(
            message,
            status_code=response.status_code,
            code=code,
            payload=payload,
        )

    async def _json(self, response: Any) -> dict[str, Any]:
        try:
            payload = response.json()
        finally:
            await response.aclose()

        if not isinstance(payload, dict):
            raise OuteAIError("Expected a JSON object in the API response.")
        return payload

    async def _data(self, response: Any) -> dict[str, Any]:
        payload = await self._json(response)
        data = payload.get("data", {})
        if not isinstance(data, dict):
            raise OuteAIError("Expected `data` to be a JSON object.")
        return data
