"""Handler for Anthropic Messages API endpoint (/v1/messages).

Enables transparent proxy mode for Claude Code and other tools that
use the Anthropic Messages format natively.

Two modes:
- **Passthrough** (model is a Claude model): Forward raw request to
  Anthropic with the caller's own ``x-api-key``. Byte-level transparent.
- **Routing** (model="auto"): Run the routing pipeline, select the
  cheapest capable model, convert formats if cross-provider.
"""

from __future__ import annotations

import json
import time

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

from robot_resources.proxy.handlers.completions import (
    _extract_prompt as _extract_prompt_openai,
)
from robot_resources.proxy.handlers.completions import (
    _extract_request_keys,
    _get_available_models,
    _get_personal_baseline,
    _get_provider_with_request_keys,
    _schedule_tracking,
    get_provider,
)
from robot_resources.proxy.models import (
    ChatCompletionRequest,
    Message,
    Usage,
)
from robot_resources.routing import async_route_prompt

router = APIRouter()
logger = structlog.get_logger()

# Anthropic API base URL for passthrough
_ANTHROPIC_BASE = "https://api.anthropic.com"

# Headers that must be forwarded to Anthropic
_FORWARD_HEADERS = {"anthropic-version", "anthropic-beta"}

# Default max_tokens when Anthropic request doesn't specify one
_DEFAULT_MAX_TOKENS = 4096


# ------------------------------------------------------------------
# Helpers
# ------------------------------------------------------------------


def _extract_prompt_from_messages(messages: list[dict]) -> str:
    """Extract prompt text from Anthropic-format messages for routing."""
    user_messages = [m.get("content", "") for m in messages if m.get("role") == "user"]
    if user_messages:
        last = user_messages[-1]
        # Content can be string or list of content blocks
        if isinstance(last, list):
            return " ".join(
                block.get("text", "") for block in last if block.get("type") == "text"
            )
        return str(last)
    return ""


def _build_forwarded_headers(
    api_key: str, raw_request: Request
) -> dict[str, str]:
    """Build headers for forwarding to Anthropic API."""
    headers = {
        "x-api-key": api_key,
        "content-type": "application/json",
    }
    for hdr in _FORWARD_HEADERS:
        val = raw_request.headers.get(hdr)
        if val:
            headers[hdr] = val
    # Default anthropic-version if not provided
    if "anthropic-version" not in headers:
        headers["anthropic-version"] = "2023-06-01"
    return headers


def _should_route_messages(model: str) -> bool:
    """Check if this messages request should be routed."""
    return model.lower() == "auto" or model.lower().startswith("rr-")


def _anthropic_messages_to_internal(body: dict) -> tuple[ChatCompletionRequest, str]:
    """Convert Anthropic Messages request to internal ChatCompletionRequest.

    Returns (request, prompt_text).
    """
    messages: list[Message] = []

    # System message
    system = body.get("system")
    if system:
        if isinstance(system, str):
            messages.append(Message(role="system", content=system))
        elif isinstance(system, list):
            text = " ".join(
                block.get("text", "") for block in system if block.get("type") == "text"
            )
            if text:
                messages.append(Message(role="system", content=text))

    # Conversation messages
    for msg in body.get("messages", []):
        content = msg.get("content", "")
        if isinstance(content, list):
            content = " ".join(
                block.get("text", "") for block in content if block.get("type") == "text"
            )
        messages.append(Message(role=msg["role"], content=str(content)))

    request = ChatCompletionRequest(
        model=body.get("model", "auto"),
        messages=messages,
        temperature=body.get("temperature", 1.0),
        max_tokens=body.get("max_tokens"),
        stream=body.get("stream", False),
    )
    prompt = _extract_prompt_openai(request)
    return request, prompt


# ------------------------------------------------------------------
# Passthrough: forward raw request to Anthropic
# ------------------------------------------------------------------


async def _passthrough(
    body: bytes,
    api_key: str,
    raw_request: Request,
) -> JSONResponse:
    """Forward the raw request to Anthropic and return the response."""
    headers = _build_forwarded_headers(api_key, raw_request)
    url = f"{_ANTHROPIC_BASE}/v1/messages"

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            response = await client.post(url, content=body, headers=headers)
            try:
                content = response.json()
            except (json.JSONDecodeError, UnicodeDecodeError):
                return JSONResponse(
                    content={
                        "type": "error",
                        "error": {
                            "type": "api_error",
                            "message": (
                                f"Upstream returned non-JSON response "
                                f"(HTTP {response.status_code})."
                            ),
                        },
                    },
                    status_code=502,
                    headers={"x-rr-mode": "passthrough"},
                )
            return JSONResponse(
                content=content,
                status_code=response.status_code,
                headers={"x-rr-mode": "passthrough"},
            )
    except (httpx.RequestError, httpx.TimeoutException) as e:
        logger.warning("passthrough_connection_error", error=str(e))
        return JSONResponse(
            status_code=502,
            content={
                "type": "error",
                "error": {
                    "type": "api_error",
                    "message": f"Failed to reach Anthropic API: {e}",
                },
            },
        )


async def _passthrough_stream(
    body: bytes,
    api_key: str,
    raw_request: Request,
):
    """Stream the raw request to Anthropic and proxy SSE chunks."""
    headers = _build_forwarded_headers(api_key, raw_request)
    url = f"{_ANTHROPIC_BASE}/v1/messages"

    try:
        async with httpx.AsyncClient(timeout=120.0) as client:
            async with client.stream("POST", url, content=body, headers=headers) as response:
                if response.status_code >= 400:
                    # Read error body and emit as SSE error event before closing
                    error_body = await response.aread()
                    try:
                        error_data = json.loads(error_body)
                    except (json.JSONDecodeError, UnicodeDecodeError):
                        error_data = {
                            "type": "error",
                            "error": {
                                "type": "api_error",
                                "message": f"Upstream error (HTTP {response.status_code})",
                            },
                        }
                    yield f"event: error\ndata: {json.dumps(error_data)}\n\n"
                    return
                async for line in response.aiter_lines():
                    yield f"{line}\n"
    except (httpx.RequestError, httpx.TimeoutException) as e:
        logger.warning("passthrough_stream_connection_error", error=str(e))
        error_data = {
            "type": "error",
            "error": {
                "type": "api_error",
                "message": f"Failed to reach Anthropic API: {e}",
            },
        }
        yield f"event: error\ndata: {json.dumps(error_data)}\n\n"


# ------------------------------------------------------------------
# Cross-provider streaming helper
# ------------------------------------------------------------------


async def _wrap_as_anthropic_sse(response: dict):
    """Wrap a complete Anthropic-format response as SSE events.

    Used when cross-provider routing produces a non-streaming result
    but the client requested streaming. Emits the standard Anthropic
    SSE event sequence: message_start → content_block_start →
    content_block_delta → content_block_stop → message_delta →
    message_stop.

    TRADEOFF: This is fake streaming — all text arrives in a single
    content_block_delta. Time-to-first-token equals full response time.

    Real streaming would parse OpenAI/Google chunks and re-emit as
    Anthropic SSE incrementally. Upside: true incremental delivery,
    better UX for tools like Claude Code. Downside: couples this
    endpoint to the internal chunk format of every provider's .stream()
    method — a format translation layer running on every chunk, with
    silent breakage risk if providers change their SSE format.

    This path only triggers when /v1/messages with model=auto routes
    to a non-Anthropic provider AND the client has keys for multiple
    providers. Narrow path in practice. Revisit if telemetry shows
    significant cross-provider traffic on /v1/messages.
    """
    msg_id = response.get("id", "msg_unknown")
    model = response.get("model", "unknown")
    text = ""
    for block in response.get("content", []):
        if block.get("type") == "text":
            text += block.get("text", "")
    usage = response.get("usage", {})

    # message_start
    msg_start = {
        "type": "message_start",
        "message": {
            "id": msg_id, "type": "message", "role": "assistant",
            "model": model, "content": [],
            "usage": {
                "input_tokens": usage.get("input_tokens", 0),
                "output_tokens": 0,
            },
        },
    }
    yield f"event: message_start\ndata: {json.dumps(msg_start)}\n\n"

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

    # content_block_delta
    block_delta = {
        "type": "content_block_delta", "index": 0,
        "delta": {"type": "text_delta", "text": text},
    }
    yield f"event: content_block_delta\ndata: {json.dumps(block_delta)}\n\n"

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

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

    # message_stop
    yield f"event: message_stop\ndata: {json.dumps({'type': 'message_stop'})}\n\n"


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


@router.post("/v1/messages", response_model=None)
async def create_message(raw_request: Request):
    """Create a message using the Anthropic Messages API format.

    Supports transparent proxy mode: reads ``x-api-key`` from the
    request headers and forwards to Anthropic. No stored keys needed.

    When ``model="auto"``, runs the Router's cost-optimization pipeline
    to select the cheapest capable model.
    """
    body_bytes = await raw_request.body()
    try:
        body = json.loads(body_bytes)
    except (json.JSONDecodeError, UnicodeDecodeError):
        return JSONResponse(
            status_code=400,
            content={
                "type": "error",
                "error": {
                    "type": "invalid_request_error",
                    "message": "Request body must be valid JSON.",
                },
            },
        )

    model = body.get("model", "")
    stream = body.get("stream", False)
    request_keys = _extract_request_keys(raw_request)

    logger.info(
        "messages_request",
        model=model,
        stream=stream,
        passthrough_keys=list(request_keys.keys()),
    )

    # ── Passthrough mode (direct Claude model) ────────────────────────
    if not _should_route_messages(model):
        api_key = request_keys.get("anthropic")
        if not api_key:
            # OpenClaw and other tools may send the Anthropic token as
            # Authorization: Bearer instead of x-api-key. Accept it on
            # this endpoint since /v1/messages is Anthropic-format only.
            api_key = request_keys.get("openai")
        if not api_key:
            # Fall back to stored key
            provider = get_provider("anthropic")
            if provider:
                api_key = provider.config.api_key

        if not api_key:
            return JSONResponse(
                status_code=401,
                content={
                    "type": "error",
                    "error": {
                        "type": "authentication_error",
                        "message": (
                            "No API key found. Provide x-api-key header "
                            "or configure ANTHROPIC_API_KEY."
                        ),
                    },
                },
            )

        if stream:
            return StreamingResponse(
                _passthrough_stream(body_bytes, api_key, raw_request),
                media_type="text/event-stream",
                headers={"x-rr-mode": "passthrough"},
            )

        return await _passthrough(body_bytes, api_key, raw_request)

    # ── Routing mode (model="auto") ───────────────────────────────────

    internal_request, prompt = _anthropic_messages_to_internal(body)
    available_models = _get_available_models(request_keys)
    personal_baseline = _get_personal_baseline()

    routing_decision = await async_route_prompt(
        prompt, models_db=available_models, baseline_model=personal_baseline
    )
    routed_model = routing_decision["selected_model"]
    provider_name = routing_decision["provider"]

    logger.info(
        "messages_routing_decision",
        original_model=model,
        routed_model=routed_model,
        provider=provider_name,
        task_type=routing_decision["task_type"],
        savings_percent=routing_decision["savings_percent"],
    )

    # Get provider (request keys take priority)
    provider, api_key_override = _get_provider_with_request_keys(
        provider_name, request_keys
    )

    if not provider:
        return JSONResponse(
            status_code=502,
            content={
                "type": "error",
                "error": {
                    "type": "api_error",
                    "message": (
                        f"No API key available for provider '{provider_name}'. "
                        "Provide the appropriate API key header or configure it."
                    ),
                },
            },
        )

    # If routed to Anthropic, forward in Anthropic format
    t_start = time.monotonic()

    if provider_name == "anthropic":
        # Rewrite model in body and forward
        body["model"] = routed_model
        forwarded_bytes = json.dumps(body).encode()
        key = api_key_override or provider.config.api_key

        if stream:
            return StreamingResponse(
                _passthrough_stream(forwarded_bytes, key, raw_request),
                media_type="text/event-stream",
                headers={"x-rr-mode": "routed", "x-rr-model": routed_model},
            )

        result = await _passthrough(forwarded_bytes, key, raw_request)
        latency_ms = int((time.monotonic() - t_start) * 1000)

        # Extract real usage from Anthropic response for tracking
        usage = Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0)
        if result.status_code == 200:
            try:
                resp_data = json.loads(result.body)
                resp_usage = resp_data.get("usage", {})
                input_tokens = resp_usage.get("input_tokens", 0)
                output_tokens = resp_usage.get("output_tokens", 0)
                usage = Usage(
                    prompt_tokens=input_tokens,
                    completion_tokens=output_tokens,
                    total_tokens=input_tokens + output_tokens,
                )
            except (json.JSONDecodeError, KeyError, TypeError):
                pass

        _schedule_tracking(
            routing_decision=routing_decision,
            provider_name=provider_name,
            model=routed_model,
            usage=usage,
            latency_ms=latency_ms,
            success=result.status_code == 200,
            prompt=prompt,
        )
        return result

    # If routed to non-Anthropic provider, use internal format conversion
    forwarded = ChatCompletionRequest(
        model=routed_model,
        messages=internal_request.messages,
        temperature=internal_request.temperature,
        max_tokens=internal_request.max_tokens or _DEFAULT_MAX_TOKENS,
        stream=False,  # Always non-streaming for cross-provider (format conversion)
    )

    try:
        response = await provider.complete(forwarded, api_key=api_key_override)
        latency_ms = int((time.monotonic() - t_start) * 1000)

        # Convert response to Anthropic Messages format
        anthropic_response = {
            "id": f"msg_{response.id}",
            "type": "message",
            "role": "assistant",
            "model": routed_model,
            "content": [
                {
                    "type": "text",
                    "text": response.choices[0].message.content,
                }
            ],
            "stop_reason": "end_turn",
            "usage": {
                "input_tokens": response.usage.prompt_tokens,
                "output_tokens": response.usage.completion_tokens,
            },
        }

        _schedule_tracking(
            routing_decision=routing_decision,
            provider_name=provider_name,
            model=routed_model,
            usage=response.usage,
            latency_ms=latency_ms,
            success=True,
            prompt=prompt,
        )

        if stream:
            # Client requested streaming — wrap the complete response
            # as Anthropic SSE events so the caller gets the format it expects
            return StreamingResponse(
                _wrap_as_anthropic_sse(anthropic_response),
                media_type="text/event-stream",
                headers={"x-rr-mode": "routed", "x-rr-model": routed_model},
            )

        return JSONResponse(
            content=anthropic_response,
            headers={"x-rr-mode": "routed", "x-rr-model": routed_model},
        )

    except (httpx.HTTPStatusError, httpx.RequestError) as e:
        logger.warning(
            "messages_provider_error",
            provider=provider_name,
            model=routed_model,
            error=str(e),
        )
        return JSONResponse(
            status_code=502,
            content={
                "type": "error",
                "error": {
                    "type": "api_error",
                    "message": f"Provider {provider_name} failed: {e}",
                },
            },
        )
