import io
from typing import Any, Optional

from openai import NOT_GIVEN, AsyncOpenAI, OpenAI
from pydub import AudioSegment

from ...base import BotBase


class STT(BotBase):
    model_id: str = "whisper-1"
    endpoint: Optional[str] = None
    api_key: str
    language: Optional[str] = None
    response_format: str = "verbose_json"

    class Config:
        """Configuration for this pydantic object."""

        extra = "allow"
        protected_namespaces = ()

    def __init__(self, /, **data: Any) -> None:
        super().__init__(**data)
        self.client = OpenAI(base_url=self.endpoint, api_key=self.api_key)
        self.aclient = AsyncOpenAI(base_url=self.endpoint, api_key=self.api_key)

    def _as2bytes(self, audio: AudioSegment) -> io.BytesIO:
        audio_bytes = io.BytesIO()
        audio.export(audio_bytes, format="mp3")
        audio_bytes.seek(0)
        audio_bytes.name = "buffer.mp3"
        return audio_bytes

    def infer(self, audio: AudioSegment) -> dict:
        audio_bytes = self._as2bytes(audio)

        trans = self.client.audio.transcriptions.create(
            model=self.model_id,
            file=audio_bytes,
            response_format=self.response_format,
            language=NOT_GIVEN if self.language is None else self.language,
        )
        return trans.to_dict()

    async def ainfer(self, audio: AudioSegment) -> dict:
        audio_bytes = self._as2bytes(audio)

        trans = await self.aclient.audio.transcriptions.create(
            model=self.model_id,
            file=audio_bytes,
            response_format=self.response_format,
            language=NOT_GIVEN if self.language is None else self.language,
        )
        return trans.to_dict()