"""
Google Chirp3 STT Adapter

Sends accumulated 16kHz slin16 PCM audio chunks to Google Cloud
Speech-to-Text API (V2 or V1) via REST for transcription using
the Chirp 3 model.

Features:
- Chirp 3 model — state-of-the-art Korean + multilingual accuracy
- REST-based (not streaming WebSocket) — accumulates 2-second chunks
- Supports both V2 (projects-based) and V1 (legacy) API endpoints
- Automatic API version detection from key format
- Base64 audio encoding for REST transport

API Key Format:
  V2: "project_id:api_key" — uses Speech V2 recognizers endpoint
  V1: "api_key" (plain)    — uses Speech V1 recognize endpoint

API Endpoints:
  V2: POST https://speech.googleapis.com/v2/projects/{project_id}/locations/global/recognizers/_:recognize?key={api_key}
  V1: POST https://speech.googleapis.com/v1/speech:recognize?key={api_key}
"""

from __future__ import annotations

import asyncio
import base64
import time
from typing import Any, AsyncIterable, Callable

import aiohttp

from dvgateway.audio.codec import float32_to_slin16
from dvgateway.types import AudioChunk, SttAdapter, TranscriptResult

# Audio accumulation settings
SAMPLE_RATE = 16000
CHUNK_DURATION_SEC = 2.0
CHUNK_SAMPLES = int(SAMPLE_RATE * CHUNK_DURATION_SEC)
CHUNK_BYTES = CHUNK_SAMPLES * 2  # 16-bit = 2 bytes per sample


class GoogleChirp3Adapter(SttAdapter):
    """Google Chirp3 STT adapter implementing the SttAdapter interface.

    Accumulates audio into 2-second chunks and sends them to Google
    Cloud Speech-to-Text API via REST for transcription. Supports
    both V2 (project-based) and V1 (legacy) API endpoints.

    The API key format determines which endpoint is used:
    - ``"project_id:api_key"`` -> V2 endpoint
    - ``"api_key"`` (no colon) -> V1 endpoint

    Args:
        api_key: Google Cloud API key. Use ``"project_id:api_key"``
            format for V2 API, or plain key for V1.
        language: BCP-47 language code (default: ko-KR).
        model: Speech recognition model (default: chirp_3).
    """

    def __init__(
        self,
        api_key: str = "",
        language: str = "ko-KR",
        model: str = "chirp_3",
    ) -> None:
        self._raw_api_key = api_key
        self._language = language
        self._model = model
        self._transcript_handler: Callable[[TranscriptResult], None] | None = None
        self._stopped = False
        self._task: asyncio.Task[None] | None = None

        # Parse API key to determine V1 vs V2
        self._project_id: str | None = None
        self._api_key: str = api_key
        if ":" in api_key:
            parts = api_key.split(":", 1)
            self._project_id = parts[0]
            self._api_key = parts[1]

    def on_transcript(self, handler: Callable[[TranscriptResult], None]) -> None:
        """Register a handler for transcription results."""
        self._transcript_handler = handler

    async def start_stream(self, linked_id: str, audio_stream: AsyncIterable[AudioChunk]) -> None:
        """Begin streaming audio for transcription.

        Accumulates audio samples into 2-second chunks and sends each
        chunk to the Google Speech API via REST. Transcription results
        are delivered through the registered on_transcript handler.

        Args:
            linked_id: DVGateway call/session linked ID.
            audio_stream: Async iterable of AudioChunk from DVGateway.
        """
        self._stopped = False
        self._task = asyncio.ensure_future(
            self._process_audio(linked_id, audio_stream)
        )

        # Wait for processing to complete
        try:
            await self._task
        except asyncio.CancelledError:
            pass

    async def stop(self) -> None:
        """Stop the stream and release resources."""
        self._stopped = True
        if self._task and not self._task.done():
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
            self._task = None

    async def _process_audio(self, linked_id: str, audio_stream: AsyncIterable[AudioChunk]) -> None:
        """Accumulate audio into chunks and send for recognition.

        Buffers incoming audio samples until 2 seconds of audio have
        been accumulated, then sends the chunk to the Speech API.
        Any remaining audio in the buffer is sent when the stream ends.
        """
        sample_buffer: list[float] = []

        async for chunk in audio_stream:
            if self._stopped:
                break

            sample_buffer.extend(chunk.samples)

            # When we have accumulated enough samples, send for recognition
            while len(sample_buffer) >= CHUNK_SAMPLES and not self._stopped:
                chunk_samples = sample_buffer[:CHUNK_SAMPLES]
                sample_buffer = sample_buffer[CHUNK_SAMPLES:]

                pcm_bytes = float32_to_slin16(chunk_samples)
                await self._recognize_chunk(linked_id, pcm_bytes, is_final=False)

        # Flush remaining audio
        if sample_buffer and not self._stopped:
            pcm_bytes = float32_to_slin16(sample_buffer)
            await self._recognize_chunk(linked_id, pcm_bytes, is_final=True)

    async def _recognize_chunk(self, linked_id: str, pcm_bytes: bytes, is_final: bool) -> None:
        """Send a PCM audio chunk to Google Speech API and handle the response.

        Args:
            linked_id: DVGateway call/session linked ID.
            pcm_bytes: Raw 16-bit little-endian PCM audio bytes at 16kHz.
            is_final: Whether this is the last chunk in the stream.
        """
        audio_b64 = base64.b64encode(pcm_bytes).decode("ascii")

        if self._project_id:
            result = await self._recognize_v2(audio_b64)
        else:
            result = await self._recognize_v1(audio_b64)

        if result and self._transcript_handler:
            transcript = TranscriptResult(
                linked_id=linked_id,
                text=result["text"],
                is_final=is_final or result.get("is_final", False),
                confidence=result.get("confidence"),
                language=self._language,
                timestamp_ms=time.time() * 1000,
                words=None,
                speaker=None,
                sentiment=None,
            )
            self._transcript_handler(transcript)

    async def _recognize_v2(self, audio_b64: str) -> dict[str, Any] | None:
        """Send recognition request to Google Speech V2 API.

        Args:
            audio_b64: Base64-encoded PCM audio.

        Returns:
            Dict with 'text', 'confidence', 'is_final' keys, or None if
            no transcript was returned.
        """
        url = (
            f"https://speech.googleapis.com/v2/projects/{self._project_id}"
            f"/locations/global/recognizers/_:recognize?key={self._api_key}"
        )

        request_body = {
            "config": {
                "autoDecodingConfig": {},
                "languageCodes": [self._language],
                "model": self._model,
                "features": {
                    "enableAutomaticPunctuation": True,
                },
            },
            "content": audio_b64,
        }

        return await self._send_request(url, request_body, api_version="v2")

    async def _recognize_v1(self, audio_b64: str) -> dict[str, Any] | None:
        """Send recognition request to Google Speech V1 API.

        Args:
            audio_b64: Base64-encoded PCM audio.

        Returns:
            Dict with 'text', 'confidence', 'is_final' keys, or None if
            no transcript was returned.
        """
        url = f"https://speech.googleapis.com/v1/speech:recognize?key={self._api_key}"

        request_body = {
            "config": {
                "encoding": "LINEAR16",
                "sampleRateHertz": SAMPLE_RATE,
                "languageCode": self._language,
                "model": self._model,
                "enableAutomaticPunctuation": True,
            },
            "audio": {
                "content": audio_b64,
            },
        }

        return await self._send_request(url, request_body, api_version="v1")

    async def _send_request(
        self, url: str, request_body: dict[str, Any], api_version: str
    ) -> dict[str, Any] | None:
        """Send HTTP POST request and parse the response.

        Args:
            url: Full API URL with key parameter.
            request_body: JSON request body.
            api_version: "v1" or "v2" for response parsing.

        Returns:
            Dict with 'text', 'confidence', 'is_final' keys, or None.
        """
        async with aiohttp.ClientSession() as session:
            async with session.post(
                url,
                headers={"Content-Type": "application/json"},
                json=request_body,
            ) as resp:
                if resp.status != 200:
                    err_body = await resp.text()
                    raise RuntimeError(
                        f"Google Speech {api_version.upper()} API failed "
                        f"(HTTP {resp.status}): {err_body}"
                    )

                data = await resp.json()

        return self._parse_response(data, api_version)

    def _parse_response(self, data: dict[str, Any], api_version: str) -> dict[str, Any] | None:
        """Parse Google Speech API response and extract transcript.

        Args:
            data: Parsed JSON response from the API.
            api_version: "v1" or "v2" for format differences.

        Returns:
            Dict with 'text', 'confidence', 'is_final' keys, or None
            if no transcript was returned.
        """
        results: list[dict[str, Any]] = []

        if api_version == "v2":
            results = data.get("results", [])
        else:
            results = data.get("results", [])

        if not results:
            return None

        # Take the first result
        result = results[0]
        alternatives = result.get("alternatives", [])
        if not alternatives:
            return None

        alt = alternatives[0]
        transcript = alt.get("transcript", "").strip()
        if not transcript:
            return None

        confidence = alt.get("confidence")

        # V2 uses resultEndOffset for timing, V1 uses resultEndTime
        is_final = result.get("isFinal", True)

        return {
            "text": transcript,
            "confidence": confidence,
            "is_final": is_final,
        }
