"""Anthropic Messages API endpoint."""

import asyncio
import uuid
from typing import Any

from fastapi import APIRouter, HTTPException
from loguru import logger
from sse_starlette.sse import EventSourceResponse

from mlx_manager.mlx_server.config import get_settings
from mlx_manager.mlx_server.errors import ProblemDetail, TimeoutProblem
from mlx_manager.mlx_server.models.ir import InferenceResult, InternalRequest, TextResult
from mlx_manager.mlx_server.schemas.anthropic import (
    AnthropicMessagesRequest,
    AnthropicMessagesResponse,
)
from mlx_manager.mlx_server.services.audit import audit_service
from mlx_manager.mlx_server.services.cloud.router import get_router
from mlx_manager.mlx_server.services.formatters import AnthropicFormatter
from mlx_manager.mlx_server.services.inference import (
    generate_chat_complete_response,
    generate_chat_stream,
)
from mlx_manager.mlx_server.utils.request_helpers import (
    timeout_error_event,
    validate_model_available,
    with_inference_timeout,
)

router = APIRouter(tags=["messages"])


@router.post(
    "/messages",
    response_model=None,
    responses={
        200: {
            "description": (
                "Anthropic Messages API response. "
                "Returns a JSON object (AnthropicMessagesResponse) when stream=false, "
                "or an SSE stream of Anthropic-format events "
                "(message_start, content_block_start, content_block_delta, "
                "content_block_stop, message_delta, message_stop) when stream=true."
            )
        },
        422: {"model": ProblemDetail, "description": "Validation Error"},
        404: {"model": ProblemDetail, "description": "Model Not Found"},
        408: {"model": TimeoutProblem, "description": "Request Timeout"},
        500: {"model": ProblemDetail, "description": "Internal Server Error"},
    },
)
async def create_message(
    request: AnthropicMessagesRequest,
) -> EventSourceResponse | AnthropicMessagesResponse:
    """Create a message using Anthropic Messages API format.

    Translates Anthropic-format request to internal format, runs inference,
    and returns Anthropic-format response. Supports both streaming and
    non-streaming modes.

    Reference: https://platform.claude.com/docs/en/api/messages
    """
    request_id = f"msg_{uuid.uuid4().hex[:24]}"
    logger.info(f"Messages request: model={request.model}, stream={request.stream}")

    # Validate model is available
    request.model = validate_model_available(request.model)

    async with audit_service.track_request(
        request_id=request_id,
        model=request.model,
        endpoint="/v1/messages",
        backend_type="local",
    ) as audit_ctx:
        try:
            # Convert to internal format
            internal = AnthropicFormatter.parse_request(request)

            # Always route (router is passthrough when no rules match)
            try:
                return await _route_and_respond(internal, request)
            except Exception as e:
                logger.warning(f"Routing failed, falling back: {e}")

            # Direct local inference path
            if request.stream:
                return await _handle_streaming(request, internal)
            else:
                result = await _handle_non_streaming(request, internal)
                # Update audit context with usage
                if result.usage:
                    audit_ctx.prompt_tokens = result.usage.input_tokens
                    audit_ctx.completion_tokens = result.usage.output_tokens
                    audit_ctx.total_tokens = result.usage.input_tokens + result.usage.output_tokens
                return result

        except HTTPException:
            raise
        except Exception as e:
            logger.exception(f"Messages API error: {e}")
            raise HTTPException(status_code=500, detail=str(e))


async def _handle_non_streaming(
    request: AnthropicMessagesRequest,
    internal: Any,  # InternalRequest from protocol.py
) -> AnthropicMessagesResponse:
    """Handle non-streaming request.

    Uses the 3-layer adapter pipeline: inference returns IR TextResult,
    AnthropicFormatter converts it to an Anthropic Messages response.
    """
    settings = get_settings()
    timeout = settings.timeout_chat_seconds

    inference_result = await with_inference_timeout(
        generate_chat_complete_response(
            model_id=internal.model,
            messages=internal.messages,
            max_tokens=internal.params.max_tokens,
            temperature=internal.params.temperature,
            top_p=internal.params.top_p or 1.0,
            stop=internal.stop,
            tools=internal.tools,
            images=internal.images,
        ),
        timeout=timeout,
        description="Messages",
    )

    formatter = AnthropicFormatter(
        model_id=request.model,
        request_id=f"msg_{uuid.uuid4().hex[:24]}",
    )
    return formatter.format_complete(
        inference_result.result,
        prompt_tokens=inference_result.prompt_tokens,
        completion_tokens=inference_result.completion_tokens,
    )


async def _handle_streaming(
    request: AnthropicMessagesRequest,
    internal: Any,  # InternalRequest from protocol.py
) -> EventSourceResponse:
    """Handle streaming request with Anthropic-format SSE events.

    Uses the 3-layer adapter pipeline: inference yields IR StreamEvents,
    AnthropicFormatter converts them to Anthropic Messages API SSE events.

    Event sequence:
    - event: message_start (initial message metadata)
    - event: content_block_start (begin content block)
    - event: content_block_delta (token chunks)
    - event: content_block_stop (end content block)
    - event: message_delta (final stop_reason and usage)
    - event: message_stop (stream complete)
    """
    settings = get_settings()
    timeout = settings.timeout_chat_seconds

    async def generate_events() -> Any:
        try:
            formatter = AnthropicFormatter(
                model_id=request.model,
                request_id=f"msg_{uuid.uuid4().hex[:24]}",
            )

            # Emit Anthropic message_start + content_block_start
            for sse in formatter.stream_start():
                yield sse

            # Apply timeout to preparation (model loading, template application)
            gen = await asyncio.wait_for(
                generate_chat_stream(
                    model_id=internal.model,
                    messages=internal.messages,
                    max_tokens=internal.params.max_tokens,
                    temperature=internal.params.temperature,
                    top_p=internal.params.top_p or 1.0,
                    stop=internal.stop,
                    tools=internal.tools,
                    images=internal.images,
                ),
                timeout=timeout,
            )

            output_tokens = 0
            async for item in gen:
                if isinstance(item, TextResult):
                    # Final result — emit closing events
                    for sse in formatter.stream_end(
                        item.finish_reason,
                        tool_calls=item.tool_calls,
                        output_tokens=output_tokens,
                    ):
                        yield sse
                else:
                    # StreamEvent — emit content delta
                    output_tokens += 1
                    for sse in formatter.stream_event(item):
                        yield sse

        except TimeoutError:
            logger.warning(f"Streaming messages request timed out after {timeout}s")
            # Send error event before closing (per CONTEXT.md streaming errors)
            yield timeout_error_event(timeout)

    return EventSourceResponse(generate_events())


async def _route_and_respond(
    ir: InternalRequest,
    request: AnthropicMessagesRequest,
) -> EventSourceResponse | AnthropicMessagesResponse:
    """Route via BackendRouter and format the RoutingOutcome.

    Handles four outcome types:
    - Passthrough streaming (raw_stream): wrap as EventSourceResponse
    - Passthrough non-streaming (raw_response): return as AnthropicMessagesResponse
    - IR streaming (ir_stream): format with AnthropicFormatter
    - IR non-streaming (ir_result): format with AnthropicFormatter
    """
    settings = get_settings()
    timeout = settings.timeout_chat_seconds
    backend_router = get_router()

    outcome = await with_inference_timeout(
        backend_router.route_request(ir),
        timeout=timeout,
        description="Messages (routed)",
    )

    # Passthrough: cloud backend returned protocol-native response
    if outcome.is_passthrough:
        if outcome.raw_stream is not None:
            raw_stream = outcome.raw_stream

            async def event_generator() -> Any:
                async for event in raw_stream:
                    yield event

            return EventSourceResponse(event_generator())
        else:
            # raw_response is already in Anthropic format
            if isinstance(outcome.raw_response, AnthropicMessagesResponse):
                return outcome.raw_response
            # Dict response - validate as AnthropicMessagesResponse
            return AnthropicMessagesResponse.model_validate(outcome.raw_response)

    # IR streaming result from router
    if outcome.ir_stream is not None:
        return _format_ir_stream(outcome.ir_stream, request)

    # IR non-streaming result from router
    if outcome.ir_result is not None:
        return _format_ir_complete(outcome.ir_result, request)

    raise RuntimeError("RoutingOutcome has no result")


def _format_ir_complete(
    inference_result: InferenceResult,
    request: AnthropicMessagesRequest,
) -> AnthropicMessagesResponse:
    """Format an InferenceResult as an AnthropicMessagesResponse."""
    formatter = AnthropicFormatter(
        model_id=request.model,
        request_id=f"msg_{uuid.uuid4().hex[:24]}",
    )
    return formatter.format_complete(
        inference_result.result,
        prompt_tokens=inference_result.prompt_tokens,
        completion_tokens=inference_result.completion_tokens,
    )


def _format_ir_stream(
    ir_stream: Any,
    request: AnthropicMessagesRequest,
) -> EventSourceResponse:
    """Format an IR stream as EventSourceResponse with Anthropic SSE format."""

    async def generate_events() -> Any:
        formatter = AnthropicFormatter(
            model_id=request.model,
            request_id=f"msg_{uuid.uuid4().hex[:24]}",
        )

        # Emit Anthropic message_start + content_block_start
        for sse in formatter.stream_start():
            yield sse

        # Stream IR events and format as Anthropic content_block_delta
        output_tokens = 0
        async for item in ir_stream:
            if isinstance(item, TextResult):
                # Final result — emit closing events
                for sse in formatter.stream_end(
                    item.finish_reason,
                    tool_calls=item.tool_calls,
                    output_tokens=output_tokens,
                ):
                    yield sse
            else:
                # StreamEvent — emit content delta
                output_tokens += 1
                for sse in formatter.stream_event(item):
                    yield sse

    return EventSourceResponse(generate_events())
