from __future__ import annotations

import base64
import json
from typing import Any, Dict, Optional

from fastapi import WebSocket, WebSocketDisconnect

from app.auth.jwt import AuthError, verify_token
from app.auth.rate_limit import QuotaError, release_voice_session, reserve_voice_session
import os

from app.providers.azure_foundry_stt import FoundrySTTProvider
from app.providers.azure_foundry_tts import FoundryTTSProvider
from app.providers.null_stt import NullSTTProvider
from app.providers.null_tts import NullTTSProvider
from app.voice.voicelive import run_voicelive_session
from app.voice.session import VoiceSession


PLAN_CONCURRENCY = {
    "free": 1,
    "pro": 3,
    "enterprise": 20,
}


async def voice_ws_endpoint(websocket: WebSocket, redis_factory) -> None:
    token = websocket.query_params.get("token")
    if not token:
        await websocket.accept()
        await websocket.send_json({"type": "error", "code": "missing_token"})
        await websocket.close(code=1008)
        return

    try:
        principal = verify_token(token, required_scope="voice:connect")
    except AuthError:
        await websocket.accept()
        await websocket.send_json({"type": "error", "code": "invalid_token"})
        await websocket.close(code=1008)
        return

    redis = await redis_factory()
    plan = principal.get("plan", "free")
    limit = PLAN_CONCURRENCY.get(plan, 1)
    sid = principal.get("sid")

    try:
        await reserve_voice_session(redis, principal.get("org"), sid, limit, ttl_seconds=300)
    except QuotaError:
        await websocket.accept()
        await websocket.send_json({"type": "error", "code": "quota_exceeded"})
        await websocket.close(code=1008)
        return

    await websocket.accept()
    engine = os.getenv("VOICE_ENGINE", "classic").lower()
    if engine == "voicelive":
        await run_voicelive_session(websocket, principal)
        return
    try:
        stt_provider = _select_stt_provider()
        tts_provider = _select_tts_provider()
    except RuntimeError as exc:
        await websocket.send_json({"type": "error", "code": "provider_config", "detail": str(exc)})
        await websocket.close(code=1011)
        return

    session = VoiceSession(websocket, principal, stt_provider=stt_provider, tts_provider=tts_provider)
    await session.start()

    try:
        while True:
            message = await websocket.receive()
            if message["type"] == "websocket.disconnect":
                break
            if "bytes" in message and message["bytes"] is not None:
                await session.on_audio(message["bytes"])
                continue
            if "text" in message and message["text"] is not None:
                await _handle_text(message["text"], session)
    except WebSocketDisconnect:
        pass
    finally:
        await session.close()
        await release_voice_session(redis, principal.get("org"), sid)


def _select_stt_provider():
    name = os.getenv("VOICE_STT_PROVIDER", "azure_foundry").lower()
    if name == "azure_foundry":
        return FoundrySTTProvider()
    return NullSTTProvider()


def _select_tts_provider():
    name = os.getenv("VOICE_TTS_PROVIDER", "azure_foundry").lower()
    if name == "azure_foundry":
        return FoundryTTSProvider()
    return NullTTSProvider()


async def _handle_text(text: str, session: VoiceSession) -> None:
    try:
        payload = json.loads(text)
    except json.JSONDecodeError:
        await session.websocket.send_json({"type": "error", "code": "invalid_json"})
        return

    event_type = payload.get("type")
    if event_type == "session.start":
        sample_rate = payload.get("sample_rate")
        if isinstance(sample_rate, int) and sample_rate > 0:
            await session.update_sample_rate(sample_rate)
        session.allow_barge_in = bool(payload.get("barge_in", False))
        await session.websocket.send_json({"type": "session.ready"})
        return
    if event_type == "session.end":
        try:
            await session.websocket.close(code=1000)
        except RuntimeError:
            return
        return
    if event_type in ("control.interrupt", "vad.start"):
        await session.interrupt(event_type)
        return
    if event_type == "audio.frame":
        audio_b64 = payload.get("audio")
        if not audio_b64:
            await session.websocket.send_json({"type": "error", "code": "missing_audio"})
            return
        try:
            audio = base64.b64decode(audio_b64)
        except Exception:
            await session.websocket.send_json({"type": "error", "code": "invalid_audio"})
            return
        await session.on_audio(audio)
        return
    if event_type == "vad.end":
        await session.end_utterance()
        return

    await session.websocket.send_json({"type": "error", "code": "unknown_event"})
