"""Khabeer (5abear) plugin for Pipecat — Saudi STT + TTS services.

    pip install pipecat-ai aiohttp
    curl -O https://speak-saudi-ai.lovable.app/integrations/khabeer_pipecat.py
    export KHABEER_API_KEY=5abear_sk_xxx

Usage:

    from khabeer_pipecat import KhabeerSTTService, KhabeerTTSService

    pipeline = Pipeline([
        transport.input(),
        KhabeerSTTService(),
        llm,
        KhabeerTTSService(voice="noura", accent="najdi"),
        transport.output(),
    ])
"""

from __future__ import annotations

import base64
import os
from typing import AsyncGenerator, Optional

import aiohttp
from loguru import logger

from pipecat.frames.frames import (
    ErrorFrame,
    Frame,
    TTSAudioRawFrame,
    TTSStartedFrame,
    TTSStoppedFrame,
    TranscriptionFrame,
)
from pipecat.services.stt_service import SegmentedSTTService
from pipecat.services.tts_service import TTSService
from pipecat.utils.time import time_now_iso8601

API_BASE = os.environ.get("KHABEER_API_BASE", "https://speak-saudi-ai.lovable.app/api/public/v1")
SAMPLE_RATE = 24000


def _wav_to_pcm(data: bytes) -> bytes:
    """Strip the RIFF header from a 16-bit PCM WAV payload."""
    idx = data.find(b"data")
    return data[idx + 8 :] if idx != -1 else data


class KhabeerTTSService(TTSService):
    """Saudi text-to-speech. Voices: noura, sara, dana, fahad, abdullah, khaled."""

    def __init__(
        self,
        *,
        api_key: Optional[str] = None,
        voice: str = "noura",
        accent: str = "aam",
        speed: float = 1.0,
        sample_rate: int = SAMPLE_RATE,
        session: Optional[aiohttp.ClientSession] = None,
        **kwargs,
    ):
        super().__init__(sample_rate=sample_rate, **kwargs)
        self._api_key = api_key or os.environ["KHABEER_API_KEY"]
        self._voice = voice
        self._accent = accent
        self._speed = speed
        self._session = session
        self._owns_session = session is None

    def can_generate_metrics(self) -> bool:
        return True

    async def _http(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession()
            self._owns_session = True
        return self._session

    async def stop(self, frame):  # pragma: no cover - lifecycle
        await super().stop(frame)
        if self._owns_session and self._session and not self._session.closed:
            await self._session.close()

    async def run_tts(self, text: str) -> AsyncGenerator[Frame, None]:
        logger.debug(f"Khabeer TTS: {text}")
        await self.start_ttfb_metrics()
        yield TTSStartedFrame()

        session = await self._http()
        try:
            async with session.post(
                f"{API_BASE}/tts-stream",
                headers={"Authorization": f"Bearer {self._api_key}"},
                json={
                    "text": text,
                    "voice": self._voice,
                    "accent": self._accent,
                    "speed": self._speed,
                },
            ) as res:
                if res.status != 200:
                    body = await res.text()
                    yield ErrorFrame(f"Khabeer TTS error {res.status}: {body}")
                    return

                await self.start_tts_usage_metrics(text)
                buffer = ""
                async for chunk in res.content.iter_any():
                    buffer += chunk.decode("utf-8", errors="ignore")
                    while "\n\n" in buffer:
                        block, buffer = buffer.split("\n\n", 1)
                        payload = None
                        event = None
                        for line in block.splitlines():
                            if line.startswith("event: "):
                                event = line[7:]
                            if line.startswith("data: "):
                                payload = line[6:]
                        if not payload or event != "audio_delta":
                            continue
                        import json

                        data = json.loads(payload)
                        if "audio" not in data:
                            continue
                        await self.stop_ttfb_metrics()
                        pcm = base64.b64decode(data["audio"])
                        yield TTSAudioRawFrame(pcm, self.sample_rate, 1)
        except Exception as exc:  # pragma: no cover - network
            yield ErrorFrame(f"Khabeer TTS exception: {exc}")
        finally:
            yield TTSStoppedFrame()


class KhabeerSTTService(SegmentedSTTService):
    """Saudi speech-to-text over segmented (VAD-delimited) audio."""

    def __init__(
        self,
        *,
        api_key: Optional[str] = None,
        accent: str = "aam",
        hints: Optional[list[str]] = None,
        session: Optional[aiohttp.ClientSession] = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self._api_key = api_key or os.environ["KHABEER_API_KEY"]
        self._accent = accent
        self._hints = hints or []
        self._session = session
        self._owns_session = session is None

    def can_generate_metrics(self) -> bool:
        return True

    async def _http(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession()
            self._owns_session = True
        return self._session

    async def stop(self, frame):  # pragma: no cover - lifecycle
        await super().stop(frame)
        if self._owns_session and self._session and not self._session.closed:
            await self._session.close()

    async def run_stt(self, audio: bytes) -> AsyncGenerator[Frame, None]:
        await self.start_processing_metrics()
        await self.start_ttfb_metrics()

        session = await self._http()
        form = aiohttp.FormData()
        form.add_field("file", audio, filename="audio.wav", content_type="audio/wav")
        form.add_field("accent", self._accent)
        if self._hints:
            form.add_field("hints", ",".join(self._hints))

        try:
            async with session.post(
                f"{API_BASE}/stt",
                headers={"Authorization": f"Bearer {self._api_key}"},
                data=form,
            ) as res:
                if res.status != 200:
                    yield ErrorFrame(f"Khabeer STT error {res.status}: {await res.text()}")
                    return
                data = await res.json()
        except Exception as exc:  # pragma: no cover - network
            yield ErrorFrame(f"Khabeer STT exception: {exc}")
            return
        finally:
            await self.stop_ttfb_metrics()
            await self.stop_processing_metrics()

        text = (data.get("text") or "").strip()
        if text:
            yield TranscriptionFrame(text, "", time_now_iso8601(), "ar-SA")