# Copyright 2025 LiveKit, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Speech-to-Text implementation for Gnani Vachana

This module provides an STT implementation that uses the Gnani Vachana API,
supporting both REST recognition and real-time streaming (WebSocket).
"""

from __future__ import annotations

import asyncio
import json
import os
from dataclasses import dataclass, replace
from typing import Any, Literal

import aiohttp

from livekit import rtc
from livekit.agents import (
    DEFAULT_API_CONNECT_OPTIONS,
    APIConnectionError,
    APIConnectOptions,
    APIStatusError,
    APITimeoutError,
    LanguageCode,
    stt,
    utils,
)
from livekit.agents.types import NOT_GIVEN, NotGivenOr
from livekit.agents.utils import AudioBuffer
from livekit.agents.utils.misc import is_given

from .log import logger

GnaniSTTFormat = Literal["verbatim", "transcribe"]

GNANI_STT_BASE_URL = "https://api.vachana.ai"

GnaniSTTLanguages = Literal[
    "bn-IN",
    "en-IN",
    "gu-IN",
    "hi-IN",
    "kn-IN",
    "ml-IN",
    "mr-IN",
    "pa-IN",
    "ta-IN",
    "te-IN",
    "en-IN,hi-IN",
]

SUPPORTED_LANGUAGES: set[str] = {
    "bn-IN",
    "en-IN",
    "gu-IN",
    "hi-IN",
    "kn-IN",
    "ml-IN",
    "mr-IN",
    "pa-IN",
    "ta-IN",
    "te-IN",
    "en-IN,hi-IN",
}

STREAM_SUPPORTED_LANGUAGES: set[str] = {
    "bn-IN",
    "en-IN",
    "gu-IN",
    "hi-IN",
    "kn-IN",
    "ml-IN",
    "mr-IN",
    "pa-IN",
    "ta-IN",
    "te-IN",
}

SAMPLE_RATE_16K = 16000
SAMPLE_RATE_8K = 8000
STREAM_CHUNK_BYTES = 1024


@dataclass
class GnaniSTTOptions:
    api_key: str
    language: str
    sample_rate: int = SAMPLE_RATE_16K
    base_url: str = GNANI_STT_BASE_URL
    preferred_language: str | None = None
    format: str = "verbatim"
    itn_native_numerals: bool = False


_DEPRECATED_STT_KWARGS = frozenset(("organization_id", "user_id", "http_session"))


def _check_deprecated_args(kwargs: dict[str, Any], *, caller: str = "STT.__init__") -> None:
    """Warn about deprecated kwargs and raise on truly unknown ones."""
    for name in _DEPRECATED_STT_KWARGS:
        if name in kwargs:
            logger.warning(f"`{name}` is deprecated and no longer used")

    unknown = set(kwargs) - _DEPRECATED_STT_KWARGS
    if unknown:
        raise TypeError(
            f"{caller}() got unexpected keyword argument(s): {', '.join(sorted(unknown))}"
        )


class STT(stt.STT):
    """Gnani Vachana Speech-to-Text implementation.

    Provides speech-to-text functionality using Gnani's Vachana platform.
    Supports REST recognition and real-time streaming via WebSocket.

    Args:
        language: BCP-47 language code (e.g. "hi-IN", "en-IN").
        api_key: Gnani API key (falls back to GNANI_API_KEY env var).
        sample_rate: Audio sample rate for streaming (8000 or 16000).
        base_url: Vachana API base URL.
        preferred_language: Force single-language model for this code.
        format: "verbatim" (default) or "transcribe" (enables ITN).
        itn_native_numerals: Render digits in native script when format="transcribe".
    """

    def __init__(
        self,
        *,
        language: str = "en-IN",
        api_key: str | None = None,
        sample_rate: int = SAMPLE_RATE_16K,
        base_url: str = GNANI_STT_BASE_URL,
        preferred_language: str | None = None,
        format: GnaniSTTFormat = "verbatim",
        itn_native_numerals: bool = False,
        **kwargs: Any,
    ) -> None:
        super().__init__(
            capabilities=stt.STTCapabilities(
                streaming=True,
                interim_results=False,
                aligned_transcript=False,
            )
        )

        _check_deprecated_args(kwargs)

        self._api_key = api_key or os.environ.get("GNANI_API_KEY")
        if not self._api_key:
            raise ValueError(
                "Gnani API key is required. "
                "Provide it directly or set GNANI_API_KEY environment variable."
            )

        if sample_rate not in (SAMPLE_RATE_8K, SAMPLE_RATE_16K):
            raise ValueError("sample_rate must be 8000 or 16000")

        self._opts = GnaniSTTOptions(
            api_key=self._api_key,
            language=language,
            sample_rate=sample_rate,
            base_url=base_url,
            preferred_language=preferred_language,
            format=format,
            itn_native_numerals=itn_native_numerals,
        )
        self._session: aiohttp.ClientSession | None = None

    @property
    def model(self) -> str:
        return "vachana-stt-v3"

    @property
    def provider(self) -> str:
        return "Gnani"

    def _ensure_session(self) -> aiohttp.ClientSession:
        if not self._session:
            self._session = utils.http_context.http_session()
        return self._session

    @staticmethod
    def _single_attempt(conn_options: APIConnectOptions) -> APIConnectOptions:
        return APIConnectOptions(
            max_retry=0,
            retry_interval=conn_options.retry_interval,
            timeout=conn_options.timeout,
        )

    async def recognize(
        self,
        buffer: AudioBuffer,
        *,
        language: NotGivenOr[str] = NOT_GIVEN,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> stt.SpeechEvent:
        return await super().recognize(
            buffer,
            language=language,
            conn_options=self._single_attempt(conn_options),
        )

    async def _recognize_impl(
        self,
        buffer: AudioBuffer,
        *,
        language: NotGivenOr[str] = NOT_GIVEN,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> stt.SpeechEvent:
        lang = language if is_given(language) else self._opts.language

        wav_bytes = rtc.combine_audio_frames(buffer).to_wav_bytes()

        form_data = aiohttp.FormData()
        form_data.add_field("audio_file", wav_bytes, filename="audio.wav", content_type="audio/wav")
        form_data.add_field("language_code", lang)
        form_data.add_field("format", self._opts.format)

        if self._opts.preferred_language is not None:
            form_data.add_field("preferred_language", self._opts.preferred_language)
        if self._opts.itn_native_numerals:
            form_data.add_field("itn_native_numerals", "true")

        headers: dict[str, str] = {
            "X-API-Key-ID": self._opts.api_key,
        }

        try:
            async with self._ensure_session().post(
                url=f"{self._opts.base_url}/stt/v3",
                data=form_data,
                headers=headers,
                timeout=aiohttp.ClientTimeout(
                    total=conn_options.timeout,
                    sock_connect=conn_options.timeout,
                ),
            ) as res:
                if res.status != 200:
                    error_text = await res.text()
                    logger.error(f"Gnani STT API error: {res.status} - {error_text}")
                    raise APIStatusError(
                        message=f"Gnani STT API Error ({res.status}): {error_text}",
                        status_code=res.status,
                        body=error_text,
                    )

                response_json = await res.json()
                transcript = response_json.get("transcript", "")
                request_id = response_json.get("request_id", "")

                return stt.SpeechEvent(
                    type=stt.SpeechEventType.FINAL_TRANSCRIPT,
                    request_id=request_id,
                    alternatives=[
                        stt.SpeechData(
                            language=LanguageCode(lang),
                            text=transcript,
                            confidence=1.0,
                        )
                    ],
                )

        except asyncio.TimeoutError as e:
            raise APITimeoutError("Gnani STT API request timed out") from e
        except (APIStatusError, APIConnectionError, APITimeoutError):
            raise
        except Exception as e:
            raise APIConnectionError(f"Gnani STT error: {e}") from e

    def stream(
        self,
        *,
        language: NotGivenOr[str] = NOT_GIVEN,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
    ) -> SpeechStream:
        opts = replace(self._opts)
        if is_given(language):
            opts.language = language
        return SpeechStream(
            stt=self,
            opts=opts,
            conn_options=self._single_attempt(conn_options),
        )

    async def aclose(self) -> None:
        pass


class SpeechStream(stt.RecognizeStream):
    """WebSocket-based streaming STT for Gnani Vachana.

    Connects to wss://api.vachana.ai/stt/v3/stream and sends raw PCM audio
    in 1024-byte chunks (512 samples, 16-bit mono).
    """

    def __init__(
        self,
        *,
        stt: STT,
        opts: GnaniSTTOptions,
        conn_options: APIConnectOptions,
    ) -> None:
        super().__init__(
            stt=stt,
            conn_options=conn_options,
            sample_rate=opts.sample_rate,
        )
        self._opts = opts

    def _build_ws_url(self) -> str:
        base = self._opts.base_url
        if base.startswith("https://"):
            ws_base = "wss://" + base[len("https://") :]
        elif base.startswith("http://"):
            ws_base = "ws://" + base[len("http://") :]
        else:
            ws_base = "wss://" + base
        return f"{ws_base}/stt/v3/stream"

    async def _run(self) -> None:
        import websockets

        ws_url = self._build_ws_url()
        headers: dict[str, str] = {
            "x-api-key-id": self._opts.api_key,
            "lang_code": self._opts.language,
            "x-sample-rate": str(self._opts.sample_rate),
        }
        if self._opts.format != "verbatim":
            headers["x-format"] = self._opts.format
        if self._opts.preferred_language is not None:
            headers["preferred_language"] = self._opts.preferred_language
        if self._opts.itn_native_numerals:
            headers["itn_native_numerals"] = "true"

        try:
            async with websockets.connect(
                ws_url,
                additional_headers=headers,
                ping_interval=20,
                ping_timeout=20,
                close_timeout=10,
            ) as ws:
                connected_msg = await asyncio.wait_for(ws.recv(), timeout=10)
                connected_data = json.loads(connected_msg)
                if connected_data.get("type") != "connected":
                    logger.warning(f"Unexpected first message from Gnani STT: {connected_data}")

                send_task = asyncio.create_task(self._send_audio(ws), name="gnani-stt-send")
                recv_task = asyncio.create_task(self._recv_messages(ws), name="gnani-stt-recv")

                try:
                    # Wait for send to finish; if recv errors first, propagate it.
                    done, _ = await asyncio.wait(
                        [send_task, recv_task],
                        return_when=asyncio.FIRST_COMPLETED,
                    )
                    for task in done:
                        task.result()

                    if send_task.done() and not recv_task.done():
                        # All audio sent. The Gnani API has no application-level
                        # end-of-stream message, so give the server a short
                        # window to flush final transcripts before closing.
                        try:
                            await asyncio.wait_for(asyncio.shield(recv_task), timeout=1.0)
                        except asyncio.TimeoutError:
                            pass
                finally:
                    await utils.aio.gracefully_cancel(send_task, recv_task)

        except websockets.exceptions.ConnectionClosed as e:
            raise APIConnectionError(f"Gnani STT WebSocket closed unexpectedly: {e}") from e
        except asyncio.TimeoutError as e:
            raise APITimeoutError("Gnani STT WebSocket connection timed out") from e
        except (APIConnectionError, APIStatusError, APITimeoutError):
            raise
        except Exception as e:
            raise APIConnectionError(f"Gnani STT WebSocket error: {e}") from e

    async def _send_audio(self, ws: Any) -> None:
        audio_buffer = bytearray()

        async for data in self._input_ch:
            if isinstance(data, self._FlushSentinel):
                if audio_buffer:
                    await ws.send(bytes(audio_buffer))
                    audio_buffer.clear()
                continue

            frame: rtc.AudioFrame = data
            raw_pcm = frame.data.tobytes()
            audio_buffer.extend(raw_pcm)

            while len(audio_buffer) >= STREAM_CHUNK_BYTES:
                chunk = bytes(audio_buffer[:STREAM_CHUNK_BYTES])
                audio_buffer = audio_buffer[STREAM_CHUNK_BYTES:]
                await ws.send(chunk)

        if audio_buffer:
            await ws.send(bytes(audio_buffer))

    async def _recv_messages(self, ws: Any) -> None:
        try:
            async for msg in ws:
                if isinstance(msg, bytes):
                    continue

                data = json.loads(msg)
                msg_type = data.get("type", "")

                if msg_type == "transcript":
                    text = data.get("text", "")
                    if not text:
                        continue

                    self._event_ch.send_nowait(
                        stt.SpeechEvent(
                            type=stt.SpeechEventType.FINAL_TRANSCRIPT,
                            request_id=data.get("segment_id", ""),
                            alternatives=[
                                stt.SpeechData(
                                    language=LanguageCode(self._opts.language),
                                    text=text,
                                    confidence=1.0,
                                )
                            ],
                        )
                    )

                elif msg_type in ("speech_start", "vad_start"):
                    self._event_ch.send_nowait(
                        stt.SpeechEvent(
                            type=stt.SpeechEventType.START_OF_SPEECH,
                        )
                    )

                elif msg_type in ("speech_end", "vad_end"):
                    self._event_ch.send_nowait(
                        stt.SpeechEvent(
                            type=stt.SpeechEventType.END_OF_SPEECH,
                        )
                    )

                elif msg_type == "processing":
                    pass

                elif msg_type == "error":
                    error_msg = data.get("message", "Unknown error")
                    logger.error(f"Gnani STT stream error: {error_msg}")
                    raise APIStatusError(
                        message=f"Gnani STT stream error: {error_msg}",
                        status_code=500,
                        body=error_msg,
                    )

        except asyncio.CancelledError:
            raise
        except (APIStatusError, APIConnectionError, APITimeoutError):
            raise
        except Exception as e:
            raise APIConnectionError(f"Error receiving Gnani STT messages: {e}") from e