"""LLMHosts OpenAI-compatible route -- ``POST /v1/chat/completions``.

Accepts the OpenAI ChatCompletion request schema, translates to the
unified internal format, dispatches to a backend, and translates the
response back to the OpenAI wire format (both streaming and non-streaming).
"""

from __future__ import annotations

import logging
import uuid
from typing import TYPE_CHECKING, Any

from fastapi import APIRouter, Request
from fastapi.responses import JSONResponse, StreamingResponse

from llmhosts.proxy.errors import InvalidRequestError, ProxyError
from llmhosts.proxy.models import (
    ChatCompletionChoice,
    ChatCompletionChunk,
    ChatCompletionChunkChoice,
    ChatCompletionRequest,
    ChatCompletionResponse,
    ChatMessage,
    DeltaMessage,
    UnifiedMessage,
    UnifiedRequest,
    UnifiedResponse,
    Usage,
)

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from llmhosts.proxy.dispatcher import BackendDispatcher

logger = logging.getLogger(__name__)

router = APIRouter()

# ---------------------------------------------------------------------------
# Conversion helpers
# ---------------------------------------------------------------------------


def _to_unified(req: ChatCompletionRequest) -> UnifiedRequest:
    """Convert an OpenAI ChatCompletionRequest to a UnifiedRequest."""
    messages: list[UnifiedMessage] = []
    for m in req.messages:
        content: str | list[dict[str, Any]] | None = m.content
        tool_calls_raw: list[dict[str, Any]] | None = None
        if m.tool_calls:
            tool_calls_raw = [tc.model_dump() for tc in m.tool_calls]
        messages.append(
            UnifiedMessage(
                role=m.role,
                content=content,
                name=m.name,
                tool_calls=tool_calls_raw,
                tool_call_id=m.tool_call_id,
            )
        )

    stop: list[str] | None = None
    if isinstance(req.stop, str):
        stop = [req.stop]
    elif isinstance(req.stop, list):
        stop = req.stop

    tools: list[dict[str, Any]] | None = req.tools

    return UnifiedRequest(
        model=req.model,
        messages=messages,
        stream=req.stream,
        temperature=req.temperature,
        max_tokens=req.max_tokens,
        top_p=req.top_p,
        stop=stop,
        tools=tools,
        extra={
            "n": req.n,
            "frequency_penalty": req.frequency_penalty,
            "presence_penalty": req.presence_penalty,
            "user": req.user,
        },
    )


def _unified_to_openai_response(unified: UnifiedResponse, model: str, response_id: str) -> ChatCompletionResponse:
    """Convert a UnifiedResponse back to the OpenAI envelope."""
    message = ChatMessage(role="assistant", content=unified.content)
    choice = ChatCompletionChoice(index=0, message=message, finish_reason=unified.finish_reason)
    usage = Usage(
        prompt_tokens=unified.prompt_tokens,
        completion_tokens=unified.completion_tokens,
        total_tokens=unified.total_tokens,
    )
    return ChatCompletionResponse(
        id=response_id,
        model=model,
        choices=[choice],
        usage=usage,
    )


# ---------------------------------------------------------------------------
# Streaming SSE generator
# ---------------------------------------------------------------------------


async def _stream_openai_sse(
    dispatcher: BackendDispatcher,
    unified_req: UnifiedRequest,
    response_id: str,
    model: str,
) -> AsyncIterator[str]:
    """Yield OpenAI-format SSE frames: ``data: {json}\\n\\n``."""
    # First chunk: role
    first_chunk = ChatCompletionChunk(
        id=response_id,
        model=model,
        choices=[
            ChatCompletionChunkChoice(
                index=0,
                delta=DeltaMessage(role="assistant", content=""),
            )
        ],
    )
    yield f"data: {first_chunk.model_dump_json()}\n\n"

    # Content chunks
    try:
        async for token in dispatcher.dispatch_stream(unified_req):
            chunk = ChatCompletionChunk(
                id=response_id,
                model=model,
                choices=[
                    ChatCompletionChunkChoice(
                        index=0,
                        delta=DeltaMessage(content=token),
                    )
                ],
            )
            yield f"data: {chunk.model_dump_json()}\n\n"
    except ProxyError:
        # Re-raise so the exception handler catches it for logging;
        # for streaming, the client sees a truncated stream.
        raise

    # Final chunk: finish_reason
    stop_chunk = ChatCompletionChunk(
        id=response_id,
        model=model,
        choices=[
            ChatCompletionChunkChoice(
                index=0,
                delta=DeltaMessage(),
                finish_reason="stop",
            )
        ],
    )
    yield f"data: {stop_chunk.model_dump_json()}\n\n"
    yield "data: [DONE]\n\n"


# ---------------------------------------------------------------------------
# Endpoint
# ---------------------------------------------------------------------------


@router.post("/v1/chat/completions", response_model=None)
async def chat_completions(body: ChatCompletionRequest, request: Request) -> StreamingResponse | JSONResponse:
    """OpenAI-compatible chat completions endpoint.

    Supports both streaming (``stream: true``) and non-streaming modes.
    """
    if not body.messages:
        raise InvalidRequestError("'messages' must be a non-empty array")
    if not body.model:
        raise InvalidRequestError("'model' is required")

    dispatcher: BackendDispatcher = request.app.state.dispatcher
    unified_req = _to_unified(body)
    response_id = f"chatcmpl-{uuid.uuid4().hex[:24]}"

    # Classify the request for observability header.
    # The Router's PII detection is pure computation (no I/O) so calling it
    # here in addition to the dispatcher's internal call is inexpensive.
    privacy_tier = "unknown"
    if hasattr(request.app.state, "router"):
        try:
            decision = await request.app.state.router.route(unified_req)
            privacy_tier = decision.privacy_tier or "public"
        except Exception:
            logger.debug("Router classification for header failed", exc_info=True)

    if body.stream:
        generator = _stream_openai_sse(dispatcher, unified_req, response_id, body.model)
        return StreamingResponse(
            generator,
            media_type="text/event-stream",
            headers={
                "Cache-Control": "no-cache",
                "Connection": "keep-alive",
                "X-Accel-Buffering": "no",
                "X-LLMHosts-Privacy-Tier": privacy_tier,
            },
        )

    unified_resp = await dispatcher.dispatch(unified_req)
    openai_resp = _unified_to_openai_response(unified_resp, body.model, response_id)
    resp = JSONResponse(content=openai_resp.model_dump())
    resp.headers["X-LLMHosts-Privacy-Tier"] = privacy_tier
    # At-capacity cloud fallback header
    fallback_info = (unified_resp.extra or {}).get("fallback_info")
    if fallback_info:
        resp.headers["X-LLMHosts-Fallback"] = (
            f"true; original={fallback_info['original']}; cloud={fallback_info['cloud']}"
        )
    # GPU Respect: indicate when we yielded to user workloads
    if (unified_resp.extra or {}).get("gpu_yielded"):
        resp.headers["X-LLMHosts-GPU-Yielded"] = "true"
    return resp
