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
    barge_in_clear_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)),
            barge_in_clear_timeout_s=max(0.1, _env_float("XAI_POOL_BARGE_IN_CLEAR_TIMEOUT_S", 1.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
        self.total_barge_ins = 0
        self.total_barge_in_reuses = 0
        self.total_barge_in_resets = 0
        self.total_barge_in_discarded_messages = 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,
            "total_barge_ins": self.total_barge_ins,
            "total_barge_in_reuses": self.total_barge_in_reuses,
            "total_barge_in_resets": self.total_barge_in_resets,
            "total_barge_in_discarded_messages": self.total_barge_in_discarded_messages,
        }


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 provider audio until audio.done.

    The task is intentionally cancellable. Barge-in is processed by the client
    control loop, which cancels this relay first and then establishes a clean
    provider boundary with text.clear/audio.clear before the slot can be reused.
    """
    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 _clear_slot_and_wait(
    slot: UpstreamSlot,
    *,
    timeout_s: float,
) -> tuple[bool, str | None, int]:
    """Clear one xAI connection and drain stale frames until audio.clear.

    Returns ``(confirmed, audio_clear_json, discarded_message_count)``.
    Anything received before ``audio.clear`` is deliberately discarded. This is
    the critical safety boundary that prevents audio from a cancelled utterance
    from leaking into the next user turn.
    """
    ws = slot.ws
    if ws is None or not slot.healthy:
        return False, None, 0

    discarded = 0
    try:
        await ws.send_str(json.dumps({"type": "text.clear"}))
    except Exception:
        return False, None, discarded

    deadline = time.monotonic() + timeout_s
    while True:
        remaining = deadline - time.monotonic()
        if remaining <= 0:
            return False, None, discarded
        try:
            msg = await asyncio.wait_for(ws.receive(), timeout=remaining)
        except asyncio.TimeoutError:
            return False, None, discarded
        except Exception:
            return False, None, discarded

        if msg.type != WSMsgType.TEXT:
            if msg.type in {WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.ERROR}:
                return False, None, discarded
            discarded += 1
            continue

        try:
            payload = json.loads(msg.data)
        except Exception:
            discarded += 1
            continue
        msg_type = str(payload.get("type") or "")
        if msg_type == "audio.clear":
            return True, msg.data, discarded
        if msg_type in {"error", "response.error"}:
            return False, None, discarded
        # audio.delta/audio.done from the interrupted utterance and any other
        # stale provider messages are drained locally and MUST NOT reach LiveKit.
        discarded += 1


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
    relay_task: asyncio.Task[bool] | None = None
    client_receive_task: asyncio.Task[Any] | None = None
    turn_started = False

    async def cancel_relay() -> None:
        nonlocal relay_task
        if relay_task is None:
            return
        if not relay_task.done():
            relay_task.cancel()
        try:
            await relay_task
        except asyncio.CancelledError:
            pass
        except Exception:
            pass
        relay_task = None

    async def clean_current_slot(*, forward_ack: bool, count_barge_in: bool) -> bool:
        nonlocal leased_slot, turn_started
        if leased_slot is None:
            return True
        await cancel_relay()
        confirmed, clear_json, discarded = await _clear_slot_and_wait(
            leased_slot,
            timeout_s=pool.config.barge_in_clear_timeout_s,
        )
        pool.total_barge_in_discarded_messages += discarded
        if count_barge_in:
            pool.total_barge_ins += 1
        if confirmed:
            if forward_ack and clear_json is not None and not client.closed:
                await client.send_str(clear_json)
            if count_barge_in:
                pool.total_barge_in_reuses += 1
            await pool.release(leased_slot, healthy=True)
        else:
            if count_barge_in:
                pool.total_barge_in_resets += 1
            await pool.release(leased_slot, healthy=False)
        leased_slot = None
        turn_started = False
        return confirmed

    try:
        client_receive_task = asyncio.create_task(client.receive())
        while not client.closed:
            waiters: set[asyncio.Task[Any]] = {client_receive_task}
            if relay_task is not None:
                waiters.add(relay_task)
            done, _ = await asyncio.wait(waiters, return_when=asyncio.FIRST_COMPLETED)

            # Prefer an incoming control message when both sides complete in the
            # same event-loop turn. This lets text.clear win the barge-in race.
            if client_receive_task in done:
                msg = client_receive_task.result()
                client_receive_task = asyncio.create_task(client.receive())

                if msg.type in {WSMsgType.CLOSE, WSMsgType.CLOSED, WSMsgType.ERROR}:
                    break
                if msg.type != WSMsgType.TEXT:
                    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":
                    # Existing active/pending utterance => barge-in/reset. The
                    # slot is only reusable after audio.clear confirms the clean
                    # boundary. Residual audio is drained and never forwarded.
                    if leased_slot is not None:
                        was_active_turn = turn_started or relay_task is not None
                        if was_active_turn:
                            confirmed = await clean_current_slot(
                                forward_ack=True,
                                count_barge_in=True,
                            )
                            if not confirmed:
                                pool.total_proxy_failures += 1
                                await client.close(code=1011, message=b"xAI cancel clear unconfirmed")
                                break
                            # A cancellation clear completes the old turn. The
                            # next synthesis will acquire a (possibly different)
                            # warm slot when the client sends its next text.clear.
                            continue

                        # Defensive repeated clear before a turn starts: keep the
                        # same leased socket and return exactly one clear ACK.
                        confirmed, clear_json, discarded = await _clear_slot_and_wait(
                            leased_slot,
                            timeout_s=pool.config.barge_in_clear_timeout_s,
                        )
                        pool.total_barge_in_discarded_messages += discarded
                        if not confirmed or clear_json is None:
                            pool.total_proxy_failures += 1
                            await pool.release(leased_slot, healthy=False)
                            leased_slot = None
                            await client.close(code=1011, message=b"xAI upstream clear failed")
                            break
                        await client.send_str(clear_json)
                        continue

                    try:
                        leased_slot = await pool.acquire()
                    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
                    confirmed, clear_json, discarded = await _clear_slot_and_wait(
                        leased_slot,
                        timeout_s=pool.config.barge_in_clear_timeout_s,
                    )
                    pool.total_barge_in_discarded_messages += discarded
                    if not confirmed or clear_json is None:
                        pool.total_proxy_failures += 1
                        await pool.release(leased_slot, healthy=False)
                        leased_slot = None
                        await client.close(code=1011, message=b"xAI upstream clear failed")
                        break
                    await client.send_str(clear_json)
                    turn_started = False
                    continue

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

                if relay_task is not None:
                    await client.send_str(json.dumps({"type": "error", "message": "text.clear required to interrupt active synthesis"}))
                    continue

                assert leased_slot.ws is not None
                try:
                    await leased_slot.ws.send_str(msg.data)
                except Exception:
                    pool.total_proxy_failures += 1
                    await pool.release(leased_slot, healthy=False)
                    leased_slot = None
                    await client.close(code=1011, message=b"xAI upstream write failure")
                    break

                if msg_type in {"text.delta", "text.done"}:
                    turn_started = True
                if msg_type == "text.done":
                    relay_task = asyncio.create_task(
                        _relay_upstream_until_boundary(client, leased_slot),
                        name=f"xai-pool-relay-{leased_slot.slot_id}",
                    )
                continue

            if relay_task is not None and relay_task in done:
                try:
                    slot_healthy = bool(relay_task.result())
                except asyncio.CancelledError:
                    # Cancellation is handled by the text.clear path above.
                    continue
                except Exception:
                    slot_healthy = False
                relay_task = None
                if leased_slot is not None:
                    await pool.release(leased_slot, healthy=slot_healthy)
                    leased_slot = None
                turn_started = False
                if not slot_healthy:
                    pool.total_proxy_failures += 1
                    await client.close(code=1011, message=b"xAI upstream failed")
                    break
    finally:
        if client_receive_task is not None and not client_receive_task.done():
            client_receive_task.cancel()
            try:
                await client_receive_task
            except asyncio.CancelledError:
                pass
        if leased_slot is not None:
            # Client disappeared during an utterance. Try the same clean-boundary
            # protocol used for barge-in so the warm socket can still be reused.
            confirmed = await clean_current_slot(forward_ack=False, count_barge_in=turn_started)
            if not confirmed:
                logger.info("XAI_POOL_SLOT_RESET_AFTER_CLIENT_DISCONNECT region=%s", pool.config.region)
    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)}',
        "# TYPE tia_xai_pool_barge_ins_total counter",
        f'tia_xai_pool_barge_ins_total{{region="{region}"}} {values.get("total_barge_ins", 0)}',
        f'tia_xai_pool_barge_in_reuses_total{{region="{region}"}} {values.get("total_barge_in_reuses", 0)}',
        f'tia_xai_pool_barge_in_resets_total{{region="{region}"}} {values.get("total_barge_in_resets", 0)}',
        f'tia_xai_pool_barge_in_discarded_messages_total{{region="{region}"}} {values.get("total_barge_in_discarded_messages", 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()
