"""EB1-GNN Preview Server: OpenAI-compatible aiohttp server for GNN-routed models."""

from __future__ import annotations

import asyncio
import json
import logging
import signal
import sys
import time
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from pathlib import Path
from typing import TYPE_CHECKING, Any

import aiohttp
from aiohttp import web

from ember.eb1_gnn.errors import (
    GatewayConnectionError,
    GatewayError,
    GatewayRateLimitError,
    GatewayTimeoutError,
    RequestValidationError,
)
from ember.takumi.model_pool import CLAUDE_OPUS

if TYPE_CHECKING:
    from ember.eb1_gnn.optimizations import OptimizedQwen3Embedding

    from ember.eb1_gnn.handlers import BatchRouteCollector
    from ember.eb1_gnn.routing import EmbeddedRouter
    from ember.eb1_gnn.trace_logger import TraceLogger
    from ember.takumi.eb1_gnn_conductor import EB1Conductor

logger = logging.getLogger(__name__)


# Model identifiers served by this server
MODEL_EB1_PREVIEW = "eb1-preview"
MODEL_EB1_FAST_PREVIEW = "eb1-fast-preview"
MODEL_EB1_FRONTIER_PREVIEW = "eb1-frontier-preview"
MODEL_EB1_DELTA_PREVIEW = "eb1-delta-preview"

# SWE-bench variants
MODEL_EB1_SWE_PREVIEW = "eb1-swe-preview"
MODEL_EB1_SWE_FAST_PREVIEW = "eb1-swe-fast-preview"
MODEL_EB1_SWE_PRO_PREVIEW = "eb1-swe-pro-preview"

SUPPORTED_MODELS = frozenset(
    {
        MODEL_EB1_PREVIEW,
        MODEL_EB1_FAST_PREVIEW,
        MODEL_EB1_FRONTIER_PREVIEW,
        MODEL_EB1_DELTA_PREVIEW,
        MODEL_EB1_SWE_PREVIEW,
        MODEL_EB1_SWE_FAST_PREVIEW,
        MODEL_EB1_SWE_PRO_PREVIEW,
    }
)

SWE_MODEL_PREFIX = "eb1-swe-"


@dataclass
class ServerConfig:
    """Configuration for the EB1-GNN server.

    Attributes:
        host: Host address to bind (default 0.0.0.0).
        port: Port number (default 8082).
        registry_path: Path to GNN model registry.
        embedding_device: Device for embedding model (cuda, cuda:0, cpu).
        embedding_dimension: Target embedding dimension (default 256).
        enable_delta: Whether to enable eb1-delta-preview (requires conductor).
        gateway_url: URL of Ember gateway for model execution.
        gcs_bucket: GCS bucket for request logging (None to disable).
        gcs_flush_interval: Seconds between GCS flushes.
        gcs_flush_count: Max entries before flush.
        embedding_cache_size: LRU cache size for embeddings (default 1000).
        api_keys: API keys for bearer-token auth. Accepts either a flat list
            of key strings (backward compat, no org attribution) or a dict
            mapping key -> org_id for per-customer billing. Empty disables auth.
        trace_dir: Directory for request/response traces (None uses default).
        trace_enabled: Whether to enable trace logging.
    """

    host: str = "0.0.0.0"
    port: int = 8082
    registry_path: Path = field(
        default_factory=lambda: Path("packages/embergraphrouter/registry/eb1-gnn-preview-latest")
    )
    embedding_device: str = "cuda"
    embedding_dimension: int = 256
    enable_delta: bool = True
    gateway_url: str = "http://localhost:8080"
    gcs_bucket: str | None = "eb1-logs"
    gcs_flush_interval: int = 60
    gcs_flush_count: int = 100
    embedding_cache_size: int = 1000
    api_keys: list[str] | dict[str, str] = field(default_factory=list)
    trace_dir: str | None = None
    trace_enabled: bool = True
    # Concurrency and resource bounds
    max_concurrent_requests: int = 20
    max_concurrent_embeddings: int = 4
    conductor_max_workers: int = 4
    connection_pool_limit: int = 50
    connection_pool_limit_per_host: int = 30
    # OOD behavior for eb1-delta-preview
    ood_explore_rate: float = 0.10
    ood_fallback_model: str = CLAUDE_OPUS


@dataclass
class ServerState:
    """Runtime state for the EB1-GNN server."""

    config: ServerConfig
    embedder: "OptimizedQwen3Embedding | None" = None
    router: "EmbeddedRouter | None" = None
    conductor: "EB1Conductor | None" = None
    trace_logger: "TraceLogger | None" = None
    startup_time: float = 0.0
    request_count: int = 0
    api_key_to_org: dict[str, str] = field(default_factory=dict)
    batch_collector: "BatchRouteCollector | None" = None
    session: aiohttp.ClientSession | None = None
    request_semaphore: asyncio.Semaphore | None = None
    embedding_semaphore: asyncio.Semaphore | None = None
    embedding_pool: ThreadPoolExecutor | None = None
    conductor_pool: ThreadPoolExecutor | None = None

    @property
    def is_ready(self) -> bool:
        """Return True if the server is ready to handle requests."""
        return self.embedder is not None and self.router is not None


def _resolve_request_id(request: web.Request) -> str:
    """Resolve request ID from inbound header, generating one when missing."""
    inbound_request_id = request.headers.get("X-Request-Id", "").strip()
    if inbound_request_id:
        return inbound_request_id
    return str(uuid.uuid4())[:8]


def _validate_server_config(config: ServerConfig) -> None:
    """Validate reliability-related config and fail closed on invalid values."""
    validation_errors: list[str] = []

    if config.max_concurrent_requests <= 0:
        validation_errors.append("max_concurrent_requests must be > 0")
    if config.max_concurrent_embeddings <= 0:
        validation_errors.append("max_concurrent_embeddings must be > 0")
    if config.conductor_max_workers <= 0:
        validation_errors.append("conductor_max_workers must be > 0")
    if config.connection_pool_limit <= 0:
        validation_errors.append("connection_pool_limit must be > 0")
    if config.connection_pool_limit_per_host <= 0:
        validation_errors.append("connection_pool_limit_per_host must be > 0")
    if config.connection_pool_limit_per_host > config.connection_pool_limit:
        validation_errors.append(
            "connection_pool_limit_per_host must be <= connection_pool_limit"
        )
    if not 0.0 <= config.ood_explore_rate <= 1.0:
        validation_errors.append("ood_explore_rate must be between 0.0 and 1.0")
    if not config.ood_fallback_model.strip():
        validation_errors.append("ood_fallback_model must be non-empty")

    if validation_errors:
        joined_errors = "; ".join(validation_errors)
        raise ValueError(f"Invalid EB1-GNN reliability config: {joined_errors}")


async def handle_health(request: web.Request) -> web.Response:
    """Return 200 if ready, 503 if still initializing."""
    state: ServerState = request.app["state"]

    if state.is_ready:
        return web.json_response({"status": "ok"})
    else:
        return web.json_response(
            {"status": "initializing"},
            status=503,
        )


async def handle_models(request: web.Request) -> web.Response:
    """Return available models in OpenAI format."""
    state: ServerState = request.app["state"]

    base_models = [
        (MODEL_EB1_PREVIEW, "Balanced: Claude Opus OOD fallback"),
        (MODEL_EB1_FAST_PREVIEW, "Fast: Gemini 3 Flash OOD fallback"),
        (MODEL_EB1_FRONTIER_PREVIEW, "Frontier: eb1-pro ensemble for hard queries"),
        (MODEL_EB1_SWE_PREVIEW, "SWE: Balanced cost/performance"),
        (MODEL_EB1_SWE_FAST_PREVIEW, "SWE Fast: Latency/cost optimized"),
        (MODEL_EB1_SWE_PRO_PREVIEW, "SWE Pro: Accuracy optimized ensembles"),
    ]

    models = [
        {
            "id": model_id,
            "object": "model",
            "created": int(state.startup_time),
            "owned_by": "ember",
            "permission": [],
            "root": model_id,
            "parent": None,
            "description": description,
        }
        for model_id, description in base_models
    ]

    # Add delta-preview if conductor is available
    if state.conductor:
        models.append(
            {
                "id": MODEL_EB1_DELTA_PREVIEW,
                "object": "model",
                "created": int(state.startup_time),
                "owned_by": "ember",
                "permission": [],
                "root": MODEL_EB1_DELTA_PREVIEW,
                "parent": None,
                "description": "Delta: LLM-orchestrated routing for complex queries",
            }
        )

    return web.json_response({"object": "list", "data": models})


async def handle_chat_completions(request: web.Request) -> web.Response:
    """Handle POST /v1/chat/completions."""
    state: ServerState = request.app["state"]
    request_id = _resolve_request_id(request)
    start_time = time.perf_counter()

    try:
        body = await request.json()
    except json.JSONDecodeError as e:
        return _error_response(
            message=f"Invalid JSON: {e}",
            error_type="invalid_request_error",
            status=400,
            request_id=request_id,
        )

    model = body.get("model", "")
    is_swe_model = model.startswith(SWE_MODEL_PREFIX)
    if model not in SUPPORTED_MODELS and not is_swe_model:
        available = list(SUPPORTED_MODELS) + [f"{SWE_MODEL_PREFIX}*"]
        return _error_response(
            message=f"Model '{model}' not found. Available: {', '.join(sorted(available))}",
            error_type="invalid_request_error",
            param="model",
            status=404,
            request_id=request_id,
        )

    if not state.is_ready:
        return _error_response(
            message="Server is still initializing. Please retry.",
            error_type="server_error",
            status=503,
            request_id=request_id,
        )

    messages = body.get("messages", [])
    if not messages or not isinstance(messages, list):
        return _error_response(
            message="messages must be a non-empty array",
            error_type="invalid_request_error",
            param="messages",
            status=400,
            request_id=request_id,
        )

    query = _extract_query(messages)
    if not query:
        return _error_response(
            message="No user message found in messages",
            error_type="invalid_request_error",
            param="messages",
            status=400,
            request_id=request_id,
        )

    stream = body.get("stream", False)
    state.request_count += 1

    logger.info(
        "Request %s: model=%s, query_len=%d, stream=%s",
        request_id,
        model,
        len(query),
        stream,
    )

    if state.request_semaphore is not None:
        try:
            await asyncio.wait_for(
                state.request_semaphore.acquire(),
                timeout=5.0,
            )
        except asyncio.TimeoutError:
            return _error_response(
                message="Server at capacity, retry later",
                error_type="rate_limit_error",
                status=429,
                request_id=request_id,
            )
    try:
        return await _dispatch_chat_completions(
            state=state,
            body=body,
            model=model,
            is_swe_model=is_swe_model,
            messages=messages,
            query=query,
            stream=stream,
            request_id=request_id,
            start_time=start_time,
            request=request,
        )
    finally:
        if state.request_semaphore is not None:
            state.request_semaphore.release()


async def _dispatch_chat_completions(
    *,
    state: ServerState,
    body: dict,
    model: str,
    is_swe_model: bool,
    messages: list[dict],
    query: str,
    stream: bool,
    request_id: str,
    start_time: float,
    request: web.Request,
) -> web.Response:
    """Dispatch a chat completions request to the appropriate handler."""
    try:
        from ember.eb1_gnn.handlers import (
            handle_eb1_delta_preview,
            handle_eb1_fast_preview,
            handle_eb1_frontier_preview,
            handle_eb1_preview,
            handle_eb1_swe_preview,
        )

        if is_swe_model:
            variant = "balanced"
            if "-fast-" in model:
                variant = "fast"
            elif "-pro-" in model:
                variant = "pro"
            response_data = await handle_eb1_swe_preview(
                state=state,
                messages=messages,
                query=query,
                request_body=body,
                request_id=request_id,
                variant=variant,
            )
        elif model == MODEL_EB1_PREVIEW:
            response_data = await handle_eb1_preview(
                state=state,
                messages=messages,
                query=query,
                request_body=body,
                request_id=request_id,
            )
        elif model == MODEL_EB1_FAST_PREVIEW:
            response_data = await handle_eb1_fast_preview(
                state=state,
                messages=messages,
                query=query,
                request_body=body,
                request_id=request_id,
            )
        elif model == MODEL_EB1_FRONTIER_PREVIEW:
            response_data = await handle_eb1_frontier_preview(
                state=state,
                messages=messages,
                query=query,
                request_body=body,
                request_id=request_id,
            )
        elif model == MODEL_EB1_DELTA_PREVIEW:
            if not state.conductor:
                return _error_response(
                    message="eb1-delta-preview is not enabled on this server",
                    error_type="invalid_request_error",
                    status=400,
                    request_id=request_id,
                )
            response_data = await handle_eb1_delta_preview(
                state=state,
                messages=messages,
                query=query,
                request_body=body,
                request_id=request_id,
            )
        else:
            return _error_response(
                message=f"Unknown model: {model}",
                error_type="invalid_request_error",
                status=400,
                request_id=request_id,
            )

        elapsed_ms = (time.perf_counter() - start_time) * 1000

        log_entry = {
            "id": response_data.get("id", f"chatcmpl-{request_id}"),
            "object": "chat.completion",
            "created": int(time.time()),
            "model": model,
            "usage": {
                "prompt_tokens": response_data.get("usage", {}).get("prompt_tokens", 0),
                "completion_tokens": response_data.get("usage", {}).get("completion_tokens", 0),
                "total_tokens": response_data.get("usage", {}).get("total_tokens", 0),
            },
            "x_ember_routing": {
                "tier": response_data.get("_tier", "unknown"),
                "actual_model": response_data.get("_selected_model", "unknown"),
                "confidence": response_data.get("_confidence", 0.0),
                "entropy": response_data.get("_entropy", 0.0),
                "policy": response_data.get("_policy", model),
                "policy_override": response_data.get("_policy_override"),
            },
            "x_ember_latency": {
                "route_ms": response_data.get("_route_time_ms", 0),
                "exec_ms": response_data.get("_exec_time_ms", 0),
                "total_ms": elapsed_ms,
            },
            "x_ember_request": {
                "request_id": request_id,
                "org_id": request.get("org_id", "anonymous"),
                "query_len": len(query),
                "stream": stream,
            },
        }
        logger.info(
            "Request %s: completed in %.1fms, tier=%s, model=%s, conf=%.3f",
            request_id,
            elapsed_ms,
            log_entry["x_ember_routing"]["tier"],
            log_entry["x_ember_routing"]["actual_model"],
            log_entry["x_ember_routing"]["confidence"],
        )
        logger.info("ANALYTICS: %s", json.dumps(log_entry))

        from ember.eb1_gnn.gcs_logger import log_request

        log_request(log_entry)

        if state.trace_logger:
            await state.trace_logger.log_request(
                request_id=request_id,
                model=model,
                endpoint="/v1/chat/completions",
                request_body=body,
                response_data=response_data,
                latency_ms=elapsed_ms,
            )

        for key in list(response_data.keys()):
            if key.startswith("_"):
                response_data.pop(key)

        if stream:
            return await _stream_response(response_data, model, request_id, request)
        else:
            return web.json_response(response_data)

    except GatewayTimeoutError as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/chat/completions"
        )
        logger.warning("Request %s: timeout: %s", request_id, e)
        return _error_response(
            message=str(e),
            error_type="timeout_error",
            status=504,
            request_id=request_id,
        )
    except GatewayRateLimitError as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/chat/completions"
        )
        logger.warning("Request %s: rate limited: %s", request_id, e)
        resp = _error_response(
            message=str(e),
            error_type="rate_limit_error",
            status=429,
            request_id=request_id,
        )
        if e.retry_after is not None:
            resp.headers["Retry-After"] = str(int(e.retry_after))
        return resp
    except GatewayConnectionError as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/chat/completions"
        )
        logger.warning("Request %s: connection error: %s", request_id, e)
        return _error_response(
            message=str(e),
            error_type="upstream_error",
            status=502,
            request_id=request_id,
        )
    except GatewayError as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/chat/completions"
        )
        logger.warning("Request %s: gateway error: %s", request_id, e)
        return _error_response(
            message=str(e),
            error_type="upstream_error",
            status=e.status,
            request_id=request_id,
        )
    except RequestValidationError as e:
        return _error_response(
            message=e.detail,
            error_type="invalid_request_error",
            status=400,
            param=e.param,
            request_id=request_id,
        )
    except Exception as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/chat/completions"
        )
        logger.exception("Request %s: failed with error: %s", request_id, e)
        return _error_response(
            message=f"Internal error: {e}",
            error_type="server_error",
            status=500,
            request_id=request_id,
        )


async def handle_responses(request: web.Request) -> web.Response:
    """Handle POST /v1/responses."""
    state: ServerState = request.app["state"]
    request_id = _resolve_request_id(request)
    start_time = time.perf_counter()

    try:
        body = await request.json()
    except json.JSONDecodeError as e:
        return _error_response(
            message=f"Invalid JSON: {e}",
            error_type="invalid_request_error",
            status=400,
            request_id=request_id,
        )

    model = body.get("model", "")
    is_swe_model = model.startswith(SWE_MODEL_PREFIX)
    if model not in SUPPORTED_MODELS and not is_swe_model:
        available = list(SUPPORTED_MODELS) + [f"{SWE_MODEL_PREFIX}*"]
        return _error_response(
            message=f"Model '{model}' not found. Available: {', '.join(sorted(available))}",
            error_type="invalid_request_error",
            param="model",
            status=404,
            request_id=request_id,
        )

    if model == MODEL_EB1_DELTA_PREVIEW:
        return _error_response(
            message=(
                "Model 'eb1-delta-preview' is only supported on "
                "/v1/chat/completions"
            ),
            error_type="invalid_request_error",
            param="model",
            code="unsupported_model_endpoint",
            status=400,
            request_id=request_id,
        )

    if not state.is_ready:
        return _error_response(
            message="Server is still initializing. Please retry.",
            error_type="server_error",
            status=503,
            request_id=request_id,
        )

    input_content = body.get("input")
    if input_content is None:
        messages = body.get("messages")
        if messages:
            input_content = messages
            body["input"] = messages
        else:
            return _error_response(
                message="input is required (or messages for Chat Completions compatibility)",
                error_type="invalid_request_error",
                param="input",
                status=400,
                request_id=request_id,
            )

    state.request_count += 1

    logger.info(
        "Request %s [responses]: model=%s",
        request_id,
        model,
    )

    if state.request_semaphore is not None:
        try:
            await asyncio.wait_for(
                state.request_semaphore.acquire(),
                timeout=5.0,
            )
        except asyncio.TimeoutError:
            return _error_response(
                message="Server at capacity, retry later",
                error_type="rate_limit_error",
                status=429,
                request_id=request_id,
            )
    try:
        return await _dispatch_responses(
            state=state,
            body=body,
            model=model,
            request_id=request_id,
            start_time=start_time,
        )
    finally:
        if state.request_semaphore is not None:
            state.request_semaphore.release()


async def _dispatch_responses(
    *,
    state: ServerState,
    body: dict,
    model: str,
    request_id: str,
    start_time: float,
) -> web.Response:
    """Dispatch a Responses API request to the handler."""
    try:
        from ember.eb1_gnn.handlers import handle_eb1_responses

        response_data = await handle_eb1_responses(
            state=state,
            payload=body,
            request_id=request_id,
        )

        elapsed_ms = (time.perf_counter() - start_time) * 1000

        logger.info(
            "Request %s [responses]: completed in %.1fms, model=%s",
            request_id,
            elapsed_ms,
            response_data.get("_selected_model", "unknown"),
        )

        if state.trace_logger:
            await state.trace_logger.log_request(
                request_id=request_id,
                model=model,
                endpoint="/v1/responses",
                request_body=body,
                response_data=response_data,
                latency_ms=elapsed_ms,
            )

        for key in list(response_data.keys()):
            if key.startswith("_"):
                response_data.pop(key)

        return web.json_response(response_data)

    except GatewayTimeoutError as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/responses"
        )
        logger.warning("Request %s [responses]: timeout: %s", request_id, e)
        return _error_response(
            message=str(e),
            error_type="timeout_error",
            status=504,
            request_id=request_id,
        )
    except GatewayRateLimitError as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/responses"
        )
        logger.warning("Request %s [responses]: rate limited: %s", request_id, e)
        resp = _error_response(
            message=str(e),
            error_type="rate_limit_error",
            status=429,
            request_id=request_id,
        )
        if e.retry_after is not None:
            resp.headers["Retry-After"] = str(int(e.retry_after))
        return resp
    except GatewayConnectionError as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/responses"
        )
        logger.warning("Request %s [responses]: connection error: %s", request_id, e)
        return _error_response(
            message=str(e),
            error_type="upstream_error",
            status=502,
            request_id=request_id,
        )
    except GatewayError as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/responses"
        )
        logger.warning("Request %s [responses]: gateway error: %s", request_id, e)
        return _error_response(
            message=str(e),
            error_type="upstream_error",
            status=e.status,
            request_id=request_id,
        )
    except RequestValidationError as e:
        return _error_response(
            message=e.detail,
            error_type="invalid_request_error",
            status=400,
            param=e.param,
            request_id=request_id,
        )
    except Exception as e:
        elapsed_ms = (time.perf_counter() - start_time) * 1000
        await _log_error_trace(
            state, request_id, model, body, elapsed_ms, e, endpoint="/v1/responses"
        )
        logger.exception(
            "Request %s [responses]: failed with error: %s",
            request_id,
            e,
        )
        return _error_response(
            message=f"Internal error: {e}",
            error_type="server_error",
            status=500,
            request_id=request_id,
        )


async def _log_error_trace(
    state: ServerState,
    request_id: str,
    model: str,
    body: dict,
    elapsed_ms: float,
    error: Exception,
    *,
    endpoint: str,
) -> None:
    """Log a failed request to the trace logger if enabled."""
    if state.trace_logger:
        await state.trace_logger.log_request(
            request_id=request_id,
            model=model,
            endpoint=endpoint,
            request_body=body,
            response_data={},
            latency_ms=elapsed_ms,
            error=str(error),
        )


def _extract_query(messages: list[dict]) -> str:
    """Extract the query text from the last user message."""
    for msg in reversed(messages):
        if msg.get("role") == "user":
            content = msg.get("content", "")
            if isinstance(content, str):
                return content
            elif isinstance(content, list):
                text_parts = [
                    part.get("text", "")
                    for part in content
                    if isinstance(part, dict) and part.get("type") == "text"
                ]
                return " ".join(text_parts)
    return ""


def _error_response(
    message: str,
    error_type: str,
    status: int,
    param: str | None = None,
    code: str | None = None,
    request_id: str | None = None,
) -> web.Response:
    """Create an OpenAI-compatible error response.

    Args:
        message: Human-readable error description.
        error_type: Error category (e.g., "invalid_request_error").
        status: HTTP status code.
        param: Name of the invalid parameter, when applicable.
        code: Machine-readable error code.
        request_id: Correlation ID included as X-Request-Id header.
    """
    error_body: dict[str, Any] = {
        "error": {
            "message": message,
            "type": error_type,
        }
    }
    if param:
        error_body["error"]["param"] = param
    if code:
        error_body["error"]["code"] = code

    headers: dict[str, str] = {}
    if request_id:
        headers["X-Request-Id"] = request_id

    return web.json_response(error_body, status=status, headers=headers)


async def _stream_response(
    response_data: dict,
    model: str,
    request_id: str,
    request: web.Request,
) -> web.StreamResponse:
    """Convert a complete response to SSE stream format."""
    response = web.StreamResponse(
        status=200,
        headers={
            "Content-Type": "text/event-stream",
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Request-Id": request_id,
        },
    )
    await response.prepare(request)

    choices = response_data.get("choices", [])
    if not choices:
        await response.write(b"data: [DONE]\n\n")
        return response

    content = choices[0].get("message", {}).get("content", "")

    chunk_id = f"chatcmpl-{request_id}"
    created = int(time.time())

    initial_chunk = {
        "id": chunk_id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "delta": {"role": "assistant", "content": ""},
                "finish_reason": None,
            }
        ],
    }
    await response.write(f"data: {json.dumps(initial_chunk)}\n\n".encode())

    content_chunk = {
        "id": chunk_id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "delta": {"content": content},
                "finish_reason": None,
            }
        ],
    }
    await response.write(f"data: {json.dumps(content_chunk)}\n\n".encode())

    final_chunk = {
        "id": chunk_id,
        "object": "chat.completion.chunk",
        "created": created,
        "model": model,
        "choices": [
            {
                "index": 0,
                "delta": {},
                "finish_reason": "stop",
            }
        ],
    }
    await response.write(f"data: {json.dumps(final_chunk)}\n\n".encode())
    await response.write(b"data: [DONE]\n\n")

    return response


async def load_models(app: web.Application) -> None:
    """Load GNN models on server startup."""
    state: ServerState = app["state"]
    config = state.config

    logger.info("Loading models from registry: %s", config.registry_path)
    start_time = time.perf_counter()

    try:
        from ember.eb1_gnn.optimizations import OptimizedQwen3Embedding

        logger.info(
            "Loading Qwen3 embedding model (dim=%d, device=%s, cache=%d)...",
            config.embedding_dimension,
            config.embedding_device,
            config.embedding_cache_size,
        )
        embed_start = time.perf_counter()
        state.embedder = OptimizedQwen3Embedding.load(
            dimension=config.embedding_dimension,
            device=config.embedding_device,
            cache_size=config.embedding_cache_size,
        )
        embed_time = time.perf_counter() - embed_start
        logger.info("Embedding model loaded in %.1fs (bfloat16)", embed_time)
        if config.embedding_cache_size > 0:
            logger.info(
                "LRU cache enabled (size=%d)",
                config.embedding_cache_size,
            )

        _validate_registry_artifacts(config.registry_path)

        from embergraphrouter.inference import TieredRouter, TieredThresholds

        from ember.eb1_gnn.routing import EmbeddedRouter

        logger.info("Loading TieredRouter from registry...")
        router_start = time.perf_counter()
        thresholds = TieredThresholds.from_registry(config.registry_path)
        raw_router = TieredRouter.from_registry(
            registry_name=config.registry_path.name,
            registry_root=config.registry_path.parent,
            thresholds=thresholds,
            device=config.embedding_device if "cuda" in config.embedding_device else "cpu",
        )
        router_time = time.perf_counter() - router_start
        logger.info("TieredRouter loaded in %.1fs", router_time)

        state.router = EmbeddedRouter(raw_router, state.embedder)
        logger.info(
            "EmbeddedRouter created with dimension validation (dim=%d)",
            state.embedder.dimension,
        )

        if state.embedding_pool is not None and state.embedding_semaphore is not None:
            from ember.eb1_gnn.handlers import BatchRouteCollector

            state.batch_collector = BatchRouteCollector(
                embedder=state.embedder,
                router=state.router,
                pool=state.embedding_pool,
                semaphore=state.embedding_semaphore,
            )
            state.batch_collector.start()
            logger.info("Batch route collector started (batch=16, window=5ms)")

        if config.enable_delta:
            try:
                from ember.takumi.eb1_gnn_conductor import EB1Conductor

                logger.info("Loading EB1Conductor...")
                conductor_start = time.perf_counter()
                state.conductor = EB1Conductor(
                    gnn_router=raw_router,  # Conductor needs raw TieredRouter
                    conductor_model=CLAUDE_OPUS,
                )
                conductor_time = time.perf_counter() - conductor_start
                logger.info("EB1Conductor loaded in %.1fs", conductor_time)
            except ImportError as e:
                logger.warning("EB1Conductor not available: %s", e)
            except Exception as e:
                logger.warning("Failed to load EB1Conductor: %s", e)

        total_time = time.perf_counter() - start_time
        logger.info(
            "All models loaded in %.1fs. Server ready.",
            total_time,
        )

    except Exception as e:
        logger.exception("Failed to load models: %s", e)
        raise


async def on_startup(app: web.Application) -> None:
    """Initialize shared resources and start background model loading."""
    state: ServerState = app["state"]
    config = state.config
    state.startup_time = time.time()

    connector = aiohttp.TCPConnector(
        limit=config.connection_pool_limit,
        limit_per_host=config.connection_pool_limit_per_host,
    )
    state.session = aiohttp.ClientSession(
        connector=connector,
        timeout=aiohttp.ClientTimeout(total=None),  # Per-request timeouts
    )
    logger.info(
        "HTTP connection pool: limit=%d, per_host=%d",
        config.connection_pool_limit,
        config.connection_pool_limit_per_host,
    )

    state.request_semaphore = asyncio.Semaphore(config.max_concurrent_requests)
    state.embedding_semaphore = asyncio.Semaphore(
        config.max_concurrent_embeddings,
    )
    logger.info(
        "Semaphores: requests=%d, embeddings=%d",
        config.max_concurrent_requests,
        config.max_concurrent_embeddings,
    )

    state.embedding_pool = ThreadPoolExecutor(
        max_workers=config.max_concurrent_embeddings,
        thread_name_prefix="embed",
    )
    state.conductor_pool = ThreadPoolExecutor(
        max_workers=config.conductor_max_workers,
        thread_name_prefix="conductor",
    )
    logger.info(
        "Thread pools: embedding=%d, conductor=%d",
        config.max_concurrent_embeddings,
        config.conductor_max_workers,
    )

    if config.trace_enabled:
        from ember.eb1_gnn.trace_logger import TraceLogger

        state.trace_logger = TraceLogger(
            trace_dir=config.trace_dir,
            enabled=True,
        )
        await state.trace_logger.initialize()
        logger.info("Trace logging enabled: %s", state.trace_logger.trace_dir)

    if config.gcs_bucket:
        from ember.eb1_gnn.gcs_logger import init_gcs_logger

        init_gcs_logger(
            bucket=config.gcs_bucket,
            enabled=True,
            flush_interval=config.gcs_flush_interval,
            flush_count=config.gcs_flush_count,
        )
        logger.info("GCS logging enabled: bucket=%s", config.gcs_bucket)

    app["model_loader"] = asyncio.create_task(load_models(app))


async def on_cleanup(app: web.Application) -> None:
    """Shut down shared resources."""
    logger.info("Shutting down EB1-GNN server...")

    state: ServerState = app["state"]

    if state.batch_collector is not None:
        await state.batch_collector.stop()
        logger.info("Batch route collector stopped")

    if state.session and not state.session.closed:
        await state.session.close()
        logger.info("HTTP session closed")

    if state.embedding_pool:
        state.embedding_pool.shutdown(wait=False)
        logger.info("Embedding thread pool shut down")
    if state.conductor_pool:
        state.conductor_pool.shutdown(wait=False)
        logger.info("Conductor thread pool shut down")

    if state.trace_logger:
        await state.trace_logger.shutdown()
        logger.info("Trace logger shutdown complete")

    from ember.eb1_gnn.gcs_logger import get_gcs_logger

    gcs_logger = get_gcs_logger()
    if gcs_logger:
        gcs_logger.stop()

    loader = app.get("model_loader")
    if loader and not loader.done():
        loader.cancel()
        try:
            await loader
        except asyncio.CancelledError:
            pass

    logger.info("Shutdown complete.")


# Paths exempt from bearer-token authentication (health probes).
_AUTH_EXEMPT_PATHS: frozenset[str] = frozenset({"/healthz", "/health"})


@web.middleware
async def auth_middleware(
    request: web.Request,
    handler: Any,
) -> web.Response:
    """Validate bearer-token auth. Attaches request["org_id"] on success."""
    state: ServerState = request.app["state"]

    if not state.api_key_to_org:
        request["org_id"] = "anonymous"
        return await handler(request)

    if request.path in _AUTH_EXEMPT_PATHS:
        return await handler(request)

    auth_header = request.headers.get("Authorization", "")
    if not auth_header.startswith("Bearer "):
        return _error_response(
            message=(
                "Missing or malformed Authorization header. "
                "Expected: Authorization: Bearer <api-key>"
            ),
            error_type="authentication_error",
            status=401,
        )

    token = auth_header[len("Bearer ") :]
    org_id = state.api_key_to_org.get(token)
    if org_id is None:
        return _error_response(
            message="Invalid API key.",
            error_type="authentication_error",
            status=401,
        )

    request["org_id"] = org_id
    return await handler(request)


_REQUIRED_REGISTRY_ARTIFACTS: tuple[str, ...] = (
    "model.pth",
    "llm_embeddings.pkl",
    "llm_descriptions.json",
    "training_graph_cache.pkl",
)


def _validate_registry_artifacts(registry_path: Path) -> None:
    """Fail fast if required registry artifacts are missing.

    Args:
        registry_path: Path to the GNN registry directory.

    Raises:
        FileNotFoundError: If registry_path doesn't exist or is missing artifacts.
    """
    if not registry_path.is_dir():
        raise FileNotFoundError(
            f"Registry path does not exist: {registry_path}. "
            f"Expected a directory containing: {', '.join(_REQUIRED_REGISTRY_ARTIFACTS)}"
        )
    missing = [
        name for name in _REQUIRED_REGISTRY_ARTIFACTS
        if not (registry_path / name).exists()
    ]
    if missing:
        raise FileNotFoundError(
            f"Registry at {registry_path} is incomplete. "
            f"Missing: {', '.join(missing)}"
        )


def _build_api_key_mapping(
    api_keys: list[str] | dict[str, str],
) -> dict[str, str]:
    """Build key -> org_id mapping from config.

    Accepts two formats:
    - ``list[str]``: flat key list (backward compat), maps each to ``"default"``.
    - ``dict[str, str]``: explicit ``{key: org_id}`` mapping.

    Args:
        api_keys: Key configuration from ServerConfig.

    Returns:
        Dict mapping each API key to its org_id.
    """
    if isinstance(api_keys, dict):
        return dict(api_keys)
    return {key: "default" for key in api_keys}


def create_app(config: ServerConfig | None = None) -> web.Application:
    """Create the aiohttp application.

    Args:
        config: Server configuration. Uses defaults if not provided.

    Returns:
        Configured aiohttp Application.
    """
    if config is None:
        config = ServerConfig()

    _validate_server_config(config)

    app = web.Application(middlewares=[auth_middleware])
    state = ServerState(config=config)
    if config.api_keys:
        state.api_key_to_org = _build_api_key_mapping(config.api_keys)
    app["state"] = state

    app.router.add_get("/healthz", handle_health)
    app.router.add_get("/health", handle_health)
    app.router.add_get("/v1/models", handle_models)
    app.router.add_post("/v1/chat/completions", handle_chat_completions)
    app.router.add_post("/v1/responses", handle_responses)

    app.on_startup.append(on_startup)
    app.on_cleanup.append(on_cleanup)

    return app


async def run_server(config: ServerConfig | None = None) -> None:
    """Run the EB1-GNN server until interrupted."""
    if config is None:
        config = ServerConfig()

    app = create_app(config)
    runner = web.AppRunner(app)
    await runner.setup()

    site = web.TCPSite(runner, host=config.host, port=config.port)
    await site.start()

    logger.info(
        "EB1-GNN server listening on http://%s:%d",
        config.host,
        config.port,
    )
    logger.info("Health check: http://%s:%d/healthz", config.host, config.port)

    loop = asyncio.get_running_loop()
    stop_event = asyncio.Event()

    def handle_signal(sig: signal.Signals) -> None:
        logger.info("Received signal %s, shutting down...", sig.name)
        stop_event.set()

    for sig in (signal.SIGTERM, signal.SIGINT):
        loop.add_signal_handler(sig, lambda s=sig: handle_signal(s))

    await stop_event.wait()
    await runner.cleanup()


def main() -> None:
    """CLI entry point for the EB1-GNN server."""
    import argparse

    parser = argparse.ArgumentParser(
        description="EB1-GNN Preview Server",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )
    parser.add_argument(
        "--host",
        type=str,
        default="0.0.0.0",
        help="Host address to bind",
    )
    parser.add_argument(
        "--port",
        type=int,
        default=8082,
        help="Port number",
    )
    parser.add_argument(
        "--registry-path",
        type=Path,
        default=Path("packages/embergraphrouter/registry/eb1-gnn-preview-latest"),
        help="Path to GNN model registry",
    )
    parser.add_argument(
        "--device",
        type=str,
        default="cuda",
        help="Device for embedding model (cuda, cuda:0, cpu)",
    )
    parser.add_argument(
        "--embedding-dim",
        type=int,
        default=256,
        help="Embedding dimension",
    )
    parser.add_argument(
        "--disable-delta",
        action="store_true",
        help="Disable eb1-delta-preview (conductor)",
    )
    parser.add_argument(
        "--gateway-url",
        type=str,
        default="http://localhost:8080",
        help="URL of Ember gateway for model execution",
    )
    parser.add_argument(
        "--ood-explore-rate",
        type=float,
        default=0.10,
        help="Fraction of OOD delta requests routed through conductor exploration [0.0, 1.0]",
    )
    parser.add_argument(
        "--ood-fallback-model",
        type=str,
        default=CLAUDE_OPUS,
        help="Safety fallback model for OOD exploit/fallback path",
    )
    parser.add_argument(
        "--log-level",
        type=str,
        default="INFO",
        choices=["DEBUG", "INFO", "WARNING", "ERROR"],
        help="Logging level",
    )
    parser.add_argument(
        "--api-keys",
        type=str,
        default="",
        help="Comma-separated API keys for bearer-token auth (empty = no auth)",
    )
    parser.add_argument(
        "--api-keys-file",
        type=Path,
        default=None,
        help='JSON file mapping API key to org_id (e.g. {"sk-abc": "org_xyz"})',
    )

    args = parser.parse_args()

    logging.basicConfig(
        level=getattr(logging, args.log_level),
        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
        stream=sys.stderr,
    )

    api_keys: list[str] | dict[str, str] = []
    if args.api_keys_file:
        with open(args.api_keys_file) as f:
            api_keys = json.load(f)
        if not isinstance(api_keys, dict):
            parser.error("--api-keys-file must contain a JSON object {key: org_id}")
        logger.info("Loaded %d API keys from %s", len(api_keys), args.api_keys_file)
    elif args.api_keys:
        api_keys = [k.strip() for k in args.api_keys.split(",") if k.strip()]

    config = ServerConfig(
        host=args.host,
        port=args.port,
        registry_path=args.registry_path,
        embedding_device=args.device,
        embedding_dimension=args.embedding_dim,
        enable_delta=not args.disable_delta,
        gateway_url=args.gateway_url,
        api_keys=api_keys,
        ood_explore_rate=args.ood_explore_rate,
        ood_fallback_model=args.ood_fallback_model,
    )

    logger.info("Starting EB1-GNN server with config: %s", config)
    asyncio.run(run_server(config))


if __name__ == "__main__":
    main()
