"""LLMHosts Anthropic-compatible route -- ``POST /v1/messages``.

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

from __future__ import annotations

import json
import logging
import uuid
from typing import TYPE_CHECKING

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

from llmhosts.proxy.errors import InvalidRequestError, ProxyError
from llmhosts.proxy.models import (
    AnthropicContentBlock,
    AnthropicRequest,
    AnthropicResponse,
    AnthropicUsage,
    UnifiedMessage,
    UnifiedRequest,
    UnifiedResponse,
)

if TYPE_CHECKING:
    from collections.abc import AsyncIterator

    from llmhosts.proxy.dispatcher import BackendDispatcher

logger = logging.getLogger(__name__)

router = APIRouter()

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


def _anthropic_content_to_text(content: str | list[AnthropicContentBlock]) -> str:
    """Flatten Anthropic content (string or block list) to plain text."""
    if isinstance(content, str):
        return content
    parts: list[str] = []
    for block in content:
        if block.text:
            parts.append(block.text)
    return "\n".join(parts)


def _to_unified(req: AnthropicRequest) -> UnifiedRequest:
    """Convert an AnthropicRequest to a UnifiedRequest."""
    messages: list[UnifiedMessage] = []

    # Anthropic's ``system`` field becomes a system message
    if req.system:
        messages.append(UnifiedMessage(role="system", content=req.system))

    for m in req.messages:
        text = _anthropic_content_to_text(m.content)
        messages.append(UnifiedMessage(role=m.role, content=text))

    stop: list[str] | None = req.stop_sequences

    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,
        extra={
            "top_k": req.top_k,
            "metadata": req.metadata.model_dump() if req.metadata else None,
        },
    )


def _unified_to_anthropic_response(unified: UnifiedResponse, model: str, response_id: str) -> AnthropicResponse:
    """Convert a UnifiedResponse back to the Anthropic envelope."""
    content_blocks: list[AnthropicContentBlock] = []
    if unified.content:
        content_blocks.append(AnthropicContentBlock(type="text", text=unified.content))

    stop_reason = "end_turn"
    if unified.finish_reason in ("stop", "end_turn"):
        stop_reason = "end_turn"
    elif unified.finish_reason == "length":
        stop_reason = "max_tokens"
    elif unified.finish_reason == "tool_use":
        stop_reason = "tool_use"

    usage = AnthropicUsage(
        input_tokens=unified.prompt_tokens,
        output_tokens=unified.completion_tokens,
    )

    return AnthropicResponse(
        id=response_id,
        type="message",
        role="assistant",
        content=content_blocks,
        model=model,
        stop_reason=stop_reason,
        usage=usage,
    )


# ---------------------------------------------------------------------------
# Streaming SSE generator (Anthropic format)
# ---------------------------------------------------------------------------


async def _stream_anthropic_sse(
    dispatcher: BackendDispatcher,
    unified_req: UnifiedRequest,
    response_id: str,
    model: str,
) -> AsyncIterator[str]:
    """Yield Anthropic-format SSE frames.

    Anthropic streaming uses ``event:`` + ``data:`` pairs::

        event: message_start
        data: {"type":"message_start","message":{...}}

        event: content_block_start
        data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}

        event: content_block_delta
        data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}

        event: content_block_stop
        data: {"type":"content_block_stop","index":0}

        event: message_delta
        data: {"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":15}}

        event: message_stop
        data: {"type":"message_stop"}
    """
    # message_start
    message_start_payload = {
        "type": "message_start",
        "message": {
            "id": response_id,
            "type": "message",
            "role": "assistant",
            "content": [],
            "model": model,
            "stop_reason": None,
            "stop_sequence": None,
            "usage": {"input_tokens": 0, "output_tokens": 0},
        },
    }
    yield f"event: message_start\ndata: {json.dumps(message_start_payload)}\n\n"

    # content_block_start
    block_start_payload = {
        "type": "content_block_start",
        "index": 0,
        "content_block": {"type": "text", "text": ""},
    }
    yield f"event: content_block_start\ndata: {json.dumps(block_start_payload)}\n\n"

    # ping (keep connection alive indicator)
    yield 'event: ping\ndata: {"type": "ping"}\n\n'

    # content_block_delta events (one per token)
    output_tokens = 0
    try:
        async for token in dispatcher.dispatch_stream(unified_req):
            output_tokens += 1  # rough approximation; real counting requires a tokenizer
            delta_payload = {
                "type": "content_block_delta",
                "index": 0,
                "delta": {"type": "text_delta", "text": token},
            }
            yield f"event: content_block_delta\ndata: {json.dumps(delta_payload)}\n\n"
    except ProxyError:
        raise

    # content_block_stop
    block_stop_payload = {"type": "content_block_stop", "index": 0}
    yield f"event: content_block_stop\ndata: {json.dumps(block_stop_payload)}\n\n"

    # message_delta
    message_delta_payload = {
        "type": "message_delta",
        "delta": {"stop_reason": "end_turn", "stop_sequence": None},
        "usage": {"output_tokens": output_tokens},
    }
    yield f"event: message_delta\ndata: {json.dumps(message_delta_payload)}\n\n"

    # message_stop
    yield 'event: message_stop\ndata: {"type": "message_stop"}\n\n'


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


@router.post("/v1/messages", response_model=None)
async def messages(body: AnthropicRequest, request: Request) -> StreamingResponse | JSONResponse:
    """Anthropic-compatible Messages 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"msg_{uuid.uuid4().hex[:24]}"

    # Classify the request for observability header.
    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_anthropic_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)
    anthropic_resp = _unified_to_anthropic_response(unified_resp, body.model, response_id)
    resp = JSONResponse(content=anthropic_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
