from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any

import requests

from .exceptions import OuteAIAPIError, OuteAIError

DEFAULT_BASE_URL = "https://outeai.com"
DEFAULT_MODEL = "outetts_1_pro"
DEFAULT_TIMEOUT = 120
DEFAULT_CHUNK_SIZE = 8192
DEFAULT_SAMPLE_RATE = 24000
DEFAULT_CHANNELS = 1
DEFAULT_SAMPLE_WIDTH = 2
SDK_VERSION = "2.0.0"


@dataclass(frozen=True)
class SpeechResult:
    wav_bytes: bytes
    content_type: str = "audio/wav"
    sample_rate: int = DEFAULT_SAMPLE_RATE
    channels: int = DEFAULT_CHANNELS
    sample_width: int = DEFAULT_SAMPLE_WIDTH

    def save(self, path: str | Path) -> Path:
        output_path = Path(path)
        output_path.write_bytes(self.wav_bytes)
        return output_path


class SpeechStream:
    def __init__(self, response: requests.Response, chunk_size: int = DEFAULT_CHUNK_SIZE) -> None:
        self._response = response
        self._chunk_size = chunk_size
        self._closed = False

    @property
    def headers(self) -> requests.structures.CaseInsensitiveDict[str]:
        return self._response.headers

    @property
    def status_code(self) -> int:
        return self._response.status_code

    def close(self) -> None:
        if not self._closed:
            self._response.close()
            self._closed = True

    def __enter__(self) -> "SpeechStream":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

    def __iter__(self):
        try:
            for chunk in self._response.iter_content(chunk_size=self._chunk_size):
                if chunk:
                    yield chunk
        finally:
            self.close()

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


class OuteAI:
    def __init__(
        self,
        api_key: str,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timeout: int | float = DEFAULT_TIMEOUT,
        session: requests.Session | 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.")

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

    def __enter__(self) -> "OuteAI":
        return self

    def __exit__(self, exc_type, exc, tb) -> None:
        self.close()

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

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

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

    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 = self._request(
            "POST",
            "/api/v1/tts_non_stream",
            json=payload,
            headers={"Accept": "audio/wav"},
        )
        try:
            wav_bytes = response.content
            content_type = response.headers.get("Content-Type", "audio/wav")
        finally:
            response.close()

        return SpeechResult(
            wav_bytes=wav_bytes,
            content_type=content_type,
        )

    def generate_speech_to_file(
        self,
        path: str | Path,
        *,
        text: str,
        voice_id: str,
        model: str = DEFAULT_MODEL,
    ) -> Path:
        result = 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,
    ) -> SpeechStream:
        payload = self._tts_payload(text=text, voice_id=voice_id, model=model)
        response = self._request(
            "POST",
            "/api/v1/tts_stream",
            json=payload,
            headers={"Accept": "audio/wav"},
            stream=True,
        )
        return SpeechStream(response=response, chunk_size=chunk_size)

    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:
        with self.stream_speech(
            text=text,
            voice_id=voice_id,
            model=model,
            chunk_size=chunk_size,
        ) as stream:
            return 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",
        }

    def _request(
        self,
        method: str,
        path: str,
        *,
        stream: bool = False,
        timeout: int | float | None = None,
        headers: dict[str, str] | None = None,
        **kwargs: Any,
    ) -> requests.Response:
        request_headers = dict(headers or {})
        try:
            response = self.session.request(
                method=method,
                url=f"{self.base_url}{path}",
                timeout=timeout or self.timeout,
                stream=stream,
                headers=request_headers,
                **kwargs,
            )
        except requests.RequestException as exc:
            raise OuteAIError(f"Request to {self.base_url}{path} failed: {exc}") from exc

        self._raise_for_error(response)
        return response

    def _raise_for_error(self, response: requests.Response) -> None:
        if response.ok:
            return

        payload: dict[str, Any] | None = None
        message = response.text.strip() or f"Request failed with status {response.status_code}."
        code = None

        try:
            payload = response.json()
        except ValueError:
            payload = None
        else:
            if isinstance(payload, dict):
                message = str(payload.get("message") or message)
                code = payload.get("code")

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

    def _json(self, response: requests.Response) -> dict[str, Any]:
        try:
            payload = response.json()
        finally:
            response.close()

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

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