from __future__ import annotations

import asyncio
import json
import logging
import os
import random
import time
from dataclasses import dataclass
from typing import Any
from urllib.parse import urlencode

import aiohttp
from aiohttp import WSMsgType, web

from app.livekit.adapters.xai_tts import (
    AUTH_METHOD_API_KEY,
    AUTH_METHOD_CONFIG_FILE,
    DEFAULT_LANGUAGE,
    DEFAULT_VOICE,
    OPTIMIZE_STREAMING_LATENCY,
    SAMPLE_RATE,
    TEXT_NORMALIZATION,
    _AuthOptions,
    _request_headers,
    _resolve_auth_method,
    _validate_config_file_auth,
)

logger = logging.getLogger("xai_pool_proxy")


def _env_int(name: str, default: int) -> int:
    try:
        return int(os.getenv(name, str(default)))
    except (TypeError, ValueError):
        return default


def _env_float(name: str, default: float) -> float:
    try:
        return float(os.getenv(name, str(default)))
    except (TypeError, ValueError):
        return default


def _env_bool(name: str, default: bool) -> bool:
    raw = os.getenv(name)
    if raw is None:
        return default
    return raw.strip().lower() in {"1", "true", "yes", "on"}


def _auth_from_env() -> _AuthOptions:
    method = _resolve_auth_method(os.getenv("XAI_TTS_AUTH_METHOD", AUTH_METHOD_API_KEY))
    api_key = os.getenv("XAI_API_KEY")
    compartment_id = os.getenv("OCI_COMPARTMENT_ID")
    config_file = os.getenv("OCI_CONFIG_FILE", "~/.oci/config")
    profile = os.getenv("OCI_CONFIG_PROFILE", "DEFAULT")

    if method == AUTH_METHOD_API_KEY:
        if not api_key:
            raise RuntimeError("XAI_API_KEY is required for xAI pool API_KEY authentication")
        return _AuthOptions(method=method, api_key=api_key)

    if not compartment_id:
        raise RuntimeError("OCI_COMPARTMENT_ID is required for IAM xAI pool authentication")
    if method == AUTH_METHOD_CONFIG_FILE:
        _validate_config_file_auth(config_file, profile)
    return _AuthOptions(
        method=method,
        compartment_id=compartment_id,
        oci_config_file=config_file if method == AUTH_METHOD_CONFIG_FILE else None,
        oci_profile=profile if method == AUTH_METHOD_CONFIG_FILE else None,
    )


@dataclass(frozen=True)
class PoolConfig:
    region: str
    upstream_url: str
    voice: str
    language: str
    size: int
    unavailable_free_threshold: int
    recover_free_threshold: int
    connect_timeout_s: float
    connection_ttl_s: float
    refresh_jitter_s: float
    maintenance_interval_s: float
    acquire_timeout_s: float

    @classmethod
    def from_env(cls) -> "PoolConfig":
        size = max(1, _env_int("XAI_POOL_SIZE", 50))
        unavailable = max(0, _env_int("XAI_POOL_UNAVAILABLE_FREE", 2))
        recover = max(unavailable + 1, _env_int("XAI_POOL_RECOVER_FREE", 5))
        recover = min(size, recover)
        return cls(
            region=(os.getenv("TIA_XAI_REGION") or "unknown").strip(),
            upstream_url=(os.getenv("XAI_POOL_UPSTREAM_URL") or os.getenv("XAI_UPSTREAM_WEBSOCKET_URL") or "").strip(),
            voice=(os.getenv("XAI_TTS_VOICE") or DEFAULT_VOICE).strip() or DEFAULT_VOICE,
            language=(os.getenv("XAI_TTS_LANGUAGE") or DEFAULT_LANGUAGE).strip() or DEFAULT_LANGUAGE,
            size=size,
            unavailable_free_threshold=min(size, unavailable),
            recover_free_threshold=recover,
            connect_timeout_s=max(0.1, _env_float("XAI_POOL_CONNECT_TIMEOUT_S", 3.0)),
            connection_ttl_s=max(30.0, _env_float("XAI_POOL_CONNECTION_TTL_S", 540.0)),
            refresh_jitter_s=max(0.0, _env_float("XAI_POOL_REFRESH_JITTER_S", 45.0)),
            maintenance_interval_s=max(0.5, _env_float("XAI_POOL_MAINTENANCE_INTERVAL_S", 2.0)),
            acquire_timeout_s=max(0.1, _env_float("XAI_POOL_ACQUIRE_TIMEOUT_S", 2.0)),
        )

    def upstream_ws_url(self) -> str:
        params = {
            "voice": self.voice,
            "language": self.language,
            "codec": "pcm",
            "sample_rate": SAMPLE_RATE,
            "optimize_streaming_latency": OPTIMIZE_STREAMING_LATENCY,
            "text_normalization": str(TEXT_NORMALIZATION).lower(),
        }
        return f"{self.upstream_url}?{urlencode(params)}"


class PoolUnavailable(RuntimeError):
    pass


class UpstreamSlot:
    def __init__(self, slot_id: int, *, config: PoolConfig, session: aiohttp.ClientSession, auth: _AuthOptions) -> None:
        self.slot_id = slot_id
        self.config = config
        self.session = session
        self.auth = auth
        self.ws: aiohttp.ClientWebSocketResponse | None = None
        self.lock = asyncio.Lock()
        self.opened_at = 0.0
        self.last_used_at = 0.0
        self.refresh_deadline = 0.0
        self.connect_failures = 0

    @property
    def leased(self) -> bool:
        return self.lock.locked()

    @property
    def healthy(self) -> bool:
        ws = self.ws
        return bool(ws is not None and not ws.closed and ws.exception() is None)

    @property
    def needs_refresh(self) -> bool:
        return bool(self.healthy and not self.leased and self.refresh_deadline > 0 and time.monotonic() >= self.refresh_deadline)

    async def connect(self) -> None:
        if self.healthy:
            return
        await self.close()
        url = self.config.upstream_ws_url()
        started = time.perf_counter()
        try:
            self.ws = await asyncio.wait_for(
                self.session.ws_connect(
                    url,
                    headers=_request_headers(self.auth, url),
                    heartbeat=None,
                    autoclose=True,
                ),
                timeout=self.config.connect_timeout_s,
            )
        except Exception:
            self.connect_failures += 1
            raise
        self.opened_at = time.monotonic()
        self.last_used_at = self.opened_at
        jitter = random.uniform(0.0, min(self.config.refresh_jitter_s, max(0.0, self.config.connection_ttl_s - 1.0)))
        self.refresh_deadline = self.opened_at + self.config.connection_ttl_s - jitter
        logger.info(
            "XAI_POOL_SLOT_OPENED slot=%s region=%s connect_ms=%s refresh_in_s=%.1f",
            self.slot_id,
            self.config.region,
            round((time.perf_counter() - started) * 1000),
            max(0.0, self.refresh_deadline - time.monotonic()),
        )

    async def close(self) -> None:
        ws, self.ws = self.ws, None
        if ws is not None and not ws.closed:
            try:
                await ws.close()
            except Exception:
                logger.debug("failed closing xAI pool slot=%s", self.slot_id, exc_info=True)

    async def refresh(self) -> None:
        if self.leased:
            return
        await self.close()
        await self.connect()


class RegionalXAIPool:
    def __init__(self, config: PoolConfig) -> None:
        if not config.upstream_url:
            raise RuntimeError("XAI_POOL_UPSTREAM_URL is required")
        self.config = config
        self.auth = _auth_from_env()
        self.session: aiohttp.ClientSession | None = None
        self.slots: list[UpstreamSlot] = []
        self._condition = asyncio.Condition()
        self._maintenance_task: asyncio.Task[None] | None = None
        self._draining = False
        self._ready = False
        self._started = False
        self.total_acquires = 0
        self.total_acquire_timeouts = 0
        self.total_proxy_failures = 0

    async def start(self) -> None:
        if self._started:
            return
        connector = aiohttp.TCPConnector(limit=max(self.config.size * 2, 100), ttl_dns_cache=300)
        self.session = aiohttp.ClientSession(connector=connector)
        self.slots = [UpstreamSlot(i + 1, config=self.config, session=self.session, auth=self.auth) for i in range(self.config.size)]
        # Prewarm in bounded waves so startup does not create one handshake burst.
        concurrency = max(1, min(self.config.size, _env_int("XAI_POOL_PREWARM_CONCURRENCY", 5)))
        sem = asyncio.Semaphore(concurrency)

        async def open_slot(slot: UpstreamSlot) -> None:
            async with sem:
                try:
                    await slot.connect()
                except Exception as exc:
                    logger.warning("XAI_POOL_PREWARM_FAILED slot=%s error=%s", slot.slot_id, type(exc).__name__)
                await asyncio.sleep(max(0.0, _env_float("XAI_POOL_PREWARM_STAGGER_S", 0.05)))

        await asyncio.gather(*(open_slot(slot) for slot in self.slots))
        self._started = True
        self._recompute_ready()
        self._maintenance_task = asyncio.create_task(self._maintenance_loop(), name="xai-pool-maintenance")
        logger.info("XAI_POOL_STARTED region=%s size=%s healthy=%s ready=%s", self.config.region, self.config.size, self.healthy_count, self._ready)

    async def stop(self) -> None:
        self._draining = True
        self._ready = False
        if self._maintenance_task is not None:
            self._maintenance_task.cancel()
            try:
                await self._maintenance_task
            except asyncio.CancelledError:
                pass
            self._maintenance_task = None
        for slot in self.slots:
            await slot.close()
        if self.session is not None:
            await self.session.close()
            self.session = None
        self._started = False

    @property
    def healthy_count(self) -> int:
        return sum(1 for slot in self.slots if slot.healthy)

    @property
    def leased_count(self) -> int:
        return sum(1 for slot in self.slots if slot.leased)

    @property
    def free_healthy_count(self) -> int:
        return sum(1 for slot in self.slots if slot.healthy and not slot.leased)

    @property
    def ready(self) -> bool:
        self._recompute_ready()
        return self._ready

    def _recompute_ready(self) -> None:
        if self._draining or not self._started:
            self._ready = False
            return
        free = self.free_healthy_count
        if self._ready:
            if free <= self.config.unavailable_free_threshold:
                self._ready = False
        else:
            if free >= self.config.recover_free_threshold:
                self._ready = True

    async def set_draining(self, draining: bool = True) -> None:
        self._draining = draining
        self._recompute_ready()
        async with self._condition:
            self._condition.notify_all()

    async def acquire(self) -> UpstreamSlot:
        deadline = time.monotonic() + self.config.acquire_timeout_s
        while True:
            if self._draining:
                raise PoolUnavailable("pool is draining")
            for slot in self.slots:
                if not slot.healthy or slot.leased:
                    continue
                if not slot.lock.locked():
                    await slot.lock.acquire()
                    if not slot.healthy:
                        slot.lock.release()
                        continue
                    self.total_acquires += 1
                    self._recompute_ready()
                    return slot
            remaining = deadline - time.monotonic()
            if remaining <= 0:
                self.total_acquire_timeouts += 1
                self._recompute_ready()
                raise PoolUnavailable("no free healthy xAI connection")
            async with self._condition:
                try:
                    await asyncio.wait_for(self._condition.wait(), timeout=min(remaining, 0.25))
                except asyncio.TimeoutError:
                    pass

    async def release(self, slot: UpstreamSlot, *, healthy: bool = True) -> None:
        slot.last_used_at = time.monotonic()
        if not healthy:
            await slot.close()
        if slot.lock.locked():
            slot.lock.release()
        self._recompute_ready()
        async with self._condition:
            self._condition.notify_all()

    async def _maintenance_loop(self) -> None:
        while True:
            await asyncio.sleep(self.config.maintenance_interval_s)
            for slot in self.slots:
                if slot.leased:
                    continue
                try:
                    if slot.needs_refresh or not slot.healthy:
                        await slot.refresh()
                except Exception as exc:
                    logger.warning("XAI_POOL_SLOT_RECOVERY_FAILED slot=%s error=%s", slot.slot_id, type(exc).__name__)
            self._recompute_ready()
            async with self._condition:
                self._condition.notify_all()

    def status(self) -> dict[str, Any]:
        self._recompute_ready()
        return {
            "status": "ready" if self._ready else "not_ready",
            "region": self.config.region,
            "draining": self._draining,
            "configured": self.config.size,
            "healthy": self.healthy_count,
            "leased": self.leased_count,
            "free": self.free_healthy_count,
            "unavailable_free_threshold": self.config.unavailable_free_threshold,
            "recover_free_threshold": self.config.recover_free_threshold,
            "total_acquires": self.total_acquires,
            "total_acquire_timeouts": self.total_acquire_timeouts,
            "total_proxy_failures": self.total_proxy_failures,
        }


POOL: RegionalXAIPool | None = None


def _client_query_matches(request: web.Request, config: PoolConfig) -> bool:
    # The pool is prewarmed for one voice/language/sample-rate profile. Explicit
    # mismatch is rejected instead of silently synthesizing with the wrong voice.
    expected = {
        "voice": config.voice,
        "language": config.language,
        "codec": "pcm",
        "sample_rate": str(SAMPLE_RATE),
    }
    for key, value in expected.items():
        incoming = request.query.get(key)
        if incoming is not None and incoming != value:
            return False
    return True


async def _relay_upstream_until_boundary(client: web.WebSocketResponse, slot: UpstreamSlot) -> bool:
    """Relay one provider response boundary. Return True only after audio.done."""
    ws = slot.ws
    if ws is None:
        return False
    while True:
        msg = await ws.receive()
        if msg.type == WSMsgType.TEXT:
            await client.send_str(msg.data)
            try:
                payload = json.loads(msg.data)
            except Exception:
                payload = {}
            if payload.get("type") == "audio.done":
                return True
            if payload.get("type") in {"error", "response.error"}:
                return False
        elif msg.type == WSMsgType.BINARY:
            await client.send_bytes(msg.data)
        elif msg.type in {WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.ERROR}:
            return False


async def websocket_proxy(request: web.Request) -> web.StreamResponse:
    pool = POOL
    if pool is None:
        raise web.HTTPServiceUnavailable(text="pool not initialized")
    if not _client_query_matches(request, pool.config):
        raise web.HTTPBadRequest(text="voice/language/codec/sample_rate differs from prewarmed pool profile")

    client = web.WebSocketResponse(heartbeat=20.0, max_msg_size=8 * 1024 * 1024)
    await client.prepare(request)
    leased_slot: UpstreamSlot | None = None
    slot_healthy = True
    try:
        async for msg in client:
            if msg.type != WSMsgType.TEXT:
                if msg.type in {WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.ERROR}:
                    break
                continue
            try:
                payload = json.loads(msg.data)
            except Exception:
                await client.send_str(json.dumps({"type": "error", "message": "invalid json"}))
                continue
            msg_type = str(payload.get("type") or "")

            if msg_type == "text.clear":
                if leased_slot is not None:
                    await pool.release(leased_slot, healthy=slot_healthy)
                    leased_slot = None
                try:
                    leased_slot = await pool.acquire()
                    slot_healthy = True
                except PoolUnavailable as exc:
                    await client.send_str(json.dumps({"type": "error", "message": str(exc), "code": "xai_pool_exhausted"}))
                    await client.close(code=1013, message=b"xAI pool exhausted")
                    break
                assert leased_slot.ws is not None
                try:
                    await leased_slot.ws.send_str(msg.data)
                    # text.clear has its own acknowledgement and must be forwarded
                    # before the client sends text.delta/text.done.
                    while True:
                        ack = await leased_slot.ws.receive()
                        if ack.type != WSMsgType.TEXT:
                            slot_healthy = False
                            raise RuntimeError("xAI clear acknowledgement failed")
                        await client.send_str(ack.data)
                        try:
                            ack_payload = json.loads(ack.data)
                        except Exception:
                            ack_payload = {}
                        ack_type = str(ack_payload.get("type") or "")
                        if ack_type == "audio.clear":
                            break
                        if ack_type in {"error", "response.error"}:
                            slot_healthy = False
                            raise RuntimeError("xAI clear returned error")
                except Exception:
                    slot_healthy = False
                    pool.total_proxy_failures += 1
                    await client.close(code=1011, message=b"xAI upstream clear failed")
                    break
                continue

            if leased_slot is None:
                await client.send_str(json.dumps({"type": "error", "message": "text.clear required before synthesis"}))
                continue

            assert leased_slot.ws is not None
            try:
                await leased_slot.ws.send_str(msg.data)
                if msg_type == "text.done":
                    slot_healthy = await _relay_upstream_until_boundary(client, leased_slot)
                    await pool.release(leased_slot, healthy=slot_healthy)
                    leased_slot = None
                    if not slot_healthy:
                        pool.total_proxy_failures += 1
                        await client.close(code=1011, message=b"xAI upstream failed")
                        break
            except Exception:
                slot_healthy = False
                pool.total_proxy_failures += 1
                await client.close(code=1011, message=b"xAI upstream failure")
                break
    finally:
        if leased_slot is not None:
            await pool.release(leased_slot, healthy=False)
    return client


async def healthz(_: web.Request) -> web.Response:
    pool = POOL
    if pool is None or not pool._started:
        return web.json_response({"status": "starting"}, status=503)
    return web.json_response({"status": "ok", "pool": pool.status()})


async def readyz(_: web.Request) -> web.Response:
    pool = POOL
    if pool is None:
        return web.json_response({"status": "not_ready", "reason": "not_initialized"}, status=503)
    payload = pool.status()
    return web.json_response(payload, status=200 if pool.ready else 503)


async def status(_: web.Request) -> web.Response:
    pool = POOL
    return web.json_response(pool.status() if pool is not None else {"status": "not_initialized"})


async def drain(_: web.Request) -> web.Response:
    pool = POOL
    if pool is not None:
        await pool.set_draining(True)
    return web.json_response({"status": "draining"})


async def metrics(_: web.Request) -> web.Response:
    pool = POOL
    values = pool.status() if pool is not None else {}
    region = str(values.get("region", "unknown")).replace('"', "")
    lines = [
        "# TYPE tia_xai_pool_connections gauge",
        f'tia_xai_pool_connections{{region="{region}",state="healthy"}} {values.get("healthy", 0)}',
        f'tia_xai_pool_connections{{region="{region}",state="leased"}} {values.get("leased", 0)}',
        f'tia_xai_pool_connections{{region="{region}",state="free"}} {values.get("free", 0)}',
        "# TYPE tia_xai_pool_acquires_total counter",
        f'tia_xai_pool_acquires_total{{region="{region}"}} {values.get("total_acquires", 0)}',
        "# TYPE tia_xai_pool_acquire_timeouts_total counter",
        f'tia_xai_pool_acquire_timeouts_total{{region="{region}"}} {values.get("total_acquire_timeouts", 0)}',
        "# TYPE tia_xai_pool_proxy_failures_total counter",
        f'tia_xai_pool_proxy_failures_total{{region="{region}"}} {values.get("total_proxy_failures", 0)}',
    ]
    return web.Response(text="\n".join(lines) + "\n", content_type="text/plain")


async def on_startup(app: web.Application) -> None:
    global POOL
    POOL = RegionalXAIPool(PoolConfig.from_env())
    await POOL.start()


async def on_cleanup(app: web.Application) -> None:
    global POOL
    if POOL is not None:
        await POOL.stop()
        POOL = None


def create_app() -> web.Application:
    app = web.Application()
    app.router.add_get("/xai/v1/tts", websocket_proxy)
    app.router.add_get("/healthz", healthz)
    app.router.add_get("/readyz", readyz)
    app.router.add_get("/pool/status", status)
    app.router.add_post("/drain", drain)
    app.router.add_get("/metrics", metrics)
    app.on_startup.append(on_startup)
    app.on_cleanup.append(on_cleanup)
    return app


def main() -> None:
    logging.basicConfig(level=os.getenv("LOG_LEVEL", "INFO"))
    web.run_app(
        create_app(),
        host=os.getenv("XAI_POOL_BIND_HOST", "0.0.0.0"),
        port=_env_int("XAI_POOL_PORT", 18100),
        access_log=logger if _env_bool("XAI_POOL_ACCESS_LOG", False) else None,
    )


if __name__ == "__main__":
    main()
