"""Main request handler — /v1/chat/completions proxy endpoint."""

from __future__ import annotations

import logging
import time

import aiohttp
from aiohttp import web

from backend.bridge.base import BridgeRequest
from backend.bridge.factory import create_bridge
from backend.config import EcateConfig, ModelConfig
from backend.language import detect_language
from backend.db.usage import UsageLogger, UsageRecord
from backend.proxy.auth import AuthContext, authenticate
from backend.proxy.fallback import CircuitBreaker
from backend.proxy.router import RouteDecision, Router

logger = logging.getLogger("ecate.proxy")


def _calculate_cost(
    model: ModelConfig,
    tokens_in: int,
    tokens_out: int,
    cache_read_tokens: int = 0,
    cache_creation_tokens: int = 0,
) -> float:
    """Calculate total cost including cache tokens.
    
    Anthropic prompt caching:
    - cache_read_tokens: charged at cache_read_cost (cheaper than input)
    - cache_creation_tokens: charged at cache_creation_cost (more expensive than input)
    - Regular input tokens are reduced by cache tokens
    """
    # Regular input tokens (excluding cached)
    regular_input = max(0, tokens_in - cache_read_tokens - cache_creation_tokens)
    
    cost = (
        regular_input * model.input_cost_per_mtok / 1_000_000
        + tokens_out * model.output_cost_per_mtok / 1_000_000
        + cache_read_tokens * model.cache_read_cost_per_mtok / 1_000_000
        + cache_creation_tokens * model.cache_creation_cost_per_mtok / 1_000_000
    )
    return cost


def _detect_images(messages: list[dict]) -> bool:
    """Check if any message contains image content."""
    for msg in messages:
        content = msg.get("content")
        if isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and part.get("type") in ("image_url", "image"):
                    return True
    return False


def _estimate_tokens(messages: list[dict]) -> int:
    """Rough token estimate from messages."""
    total = 0
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            total += len(content) // 4
        elif isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and part.get("type") == "text":
                    total += len(part.get("text", "")) // 4
    return max(1, total)


def _extract_prompt(messages: list[dict]) -> str:
    """Extract the main user prompt from messages for routing."""
    # Get 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):
                # Extract text parts
                text_parts = []
                for part in content:
                    if isinstance(part, dict) and part.get("type") == "text":
                        text_parts.append(part.get("text", ""))
                return " ".join(text_parts)
    return ""


def _parse_strategy(model_field: str) -> tuple[str, str | None]:
    """Parse the model field to extract routing strategy.

    Returns (strategy, direct_model) where direct_model is set if bypassing routing.

    Examples:
        "auto"       → ("auto", None)
        "auto:best"  → ("best", None)
        "auto:cheap" → ("cheap", None)
        "gpt-4o"     → (None, "gpt-4o")
        "openai/gpt-4o" → (None, "openai/gpt-4o")
    """
    if model_field == "auto":
        return "auto", None
    if model_field.startswith("auto:"):
        return model_field[5:], None
    return "auto", model_field  # Direct model passthrough


class ProxyHandler:
    """Handles /v1/chat/completions requests with routing, bridging, and logging."""

    def __init__(
        self,
        config: EcateConfig,
        router: Router,
        circuit_breaker: CircuitBreaker,
        usage_logger: UsageLogger,
        http_session: aiohttp.ClientSession,
    ):
        self.config = config
        self.router = router
        self.circuit_breaker = circuit_breaker
        self.usage_logger = usage_logger
        self.http_session = http_session

    def _get_fallback_decision(
        self,
        prompt: str,
        strategy: str,
        auth: AuthContext,
        needs_tools: bool,
        needs_vision: bool,
        token_estimate: int,
        exclude_models: set[str],
        language: str = "en",
    ) -> RouteDecision | None:
        """Get next best model, excluding already-tried ones."""
        decision = self.router.route(
            prompt=prompt,
            strategy=strategy or auth.strategy,
            needs_tools=needs_tools,
            needs_vision=needs_vision,
            token_estimate=token_estimate,
            allowed_models=auth.allowed_models,
            exclude_models=exclude_models,
            language=language,
        )
        return decision

    async def handle_completions(self, request: web.Request) -> web.Response:
        """Handle POST /v1/chat/completions."""
        start_time = time.monotonic()

        # Auth
        auth = await authenticate(request)

        # Parse request body
        try:
            body = await request.json()
        except Exception:
            raise web.HTTPBadRequest(text="Invalid JSON body")

        messages = body.get("messages", [])
        if not messages:
            raise web.HTTPBadRequest(text="'messages' field is required")

        model_field = body.get("model", "auto")
        stream = body.get("stream", False)
        tools = body.get("tools")

        # Determine routing strategy
        strategy, direct_model = _parse_strategy(model_field)

        # Extract prompt and compute routing inputs once
        prompt = _extract_prompt(messages)
        has_images = _detect_images(messages)
        token_estimate = _estimate_tokens(messages)
        language = detect_language(prompt) if prompt else None

        # If direct model specified, bypass routing
        if direct_model:
            resolved = self.config.get_model(direct_model)
            if not resolved:
                raise web.HTTPBadRequest(text=f"Unknown model: {direct_model}")
            provider, model = resolved
            cost = (model.input_cost_per_mtok + model.output_cost_per_mtok) / 2
            decision = RouteDecision(
                provider=provider, model=model,
                score=0.0, adjusted_score=0.0,
                cost_per_mtok=cost, strategy="direct",
            )
        else:
            # Route using embeddings
            decision = self.router.route(
                prompt=prompt,
                strategy=strategy or auth.strategy,
                needs_tools=bool(tools),
                needs_vision=has_images,
                token_estimate=token_estimate,
                allowed_models=auth.allowed_models,
                language=language,
            )

            if not decision:
                raise web.HTTPServiceUnavailable(text="No suitable model available. Run calibration first.")

        model_id = f"{decision.provider.id}/{decision.model.id}"

        # Build base bridge request
        bridge_request = BridgeRequest(
            messages=messages,
            model=decision.model.id,
            temperature=body.get("temperature"),
            max_tokens=body.get("max_tokens"),
            top_p=body.get("top_p"),
            stream=stream,
            tools=tools,
            tool_choice=body.get("tool_choice"),
            stop=body.get("stop"),
            extra={k: v for k, v in body.items()
                   if k not in ("messages", "model", "temperature", "max_tokens", "top_p",
                                "stream", "tools", "tool_choice", "stop")},
        )

        # Try primary model, then fallback if enabled
        tried_models: set[str] = set()
        max_attempts = self.config.routing.max_retries + 1 if self.config.routing.enable_fallback else 1
        last_error: Exception | None = None
        fallback_from: str | None = None

        for attempt in range(max_attempts):
            model_id = f"{decision.provider.id}/{decision.model.id}"

            # Check circuit breaker
            if not self.circuit_breaker.is_available(model_id):
                logger.warning(f"Circuit open for {model_id}, trying fallback")
                tried_models.add(model_id)
                decision = self._get_fallback_decision(
                    prompt, strategy, auth,
                    needs_tools=bool(tools),
                    needs_vision=has_images,
                    token_estimate=token_estimate,
                    exclude_models=tried_models,
                    language=language,
                )
                if not decision:
                    break
                continue

            bridge_request.model = decision.model.id
            bridge = create_bridge(decision.provider, self.http_session)

            try:
                if stream:
                    return await self._handle_stream(
                        request, bridge, bridge_request, decision, auth, start_time, model_id,
                        was_fallback=fallback_from is not None, fallback_from=fallback_from,
                    )
                else:
                    return await self._handle_complete(
                        bridge, bridge_request, decision, auth, start_time, model_id,
                        was_fallback=fallback_from is not None, fallback_from=fallback_from,
                    )
            except Exception as e:
                self.circuit_breaker.record_failure(model_id)
                logger.warning(f"Attempt {attempt + 1} failed for {model_id}: {e}")
                last_error = e
                tried_models.add(model_id)
                fallback_from = model_id

                # Try to get fallback model
                if attempt < max_attempts - 1:
                    decision = self._get_fallback_decision(
                        prompt, strategy, auth,
                        needs_tools=bool(tools),
                        needs_vision=has_images,
                        token_estimate=token_estimate,
                        exclude_models=tried_models,
                        language=language,
                    )
                    if not decision:
                        break
                    logger.info(f"Falling back to {decision.provider.id}/{decision.model.id}")

        # All attempts failed
        raise web.HTTPBadGateway(text=f"Upstream error: {str(last_error)[:200] if last_error else 'No models available'}")

    async def _handle_complete(
        self, bridge, bridge_request, decision, auth, start_time, model_id,
        was_fallback: bool = False, fallback_from: str | None = None,
    ) -> web.Response:
        response = await bridge.complete(bridge_request, decision.model)
        self.circuit_breaker.record_success(model_id)

        elapsed = (time.monotonic() - start_time) * 1000

        # Build OpenAI-format response with cache info if available
        usage_block = {
            "prompt_tokens": response.tokens_in,
            "completion_tokens": response.tokens_out,
            "total_tokens": response.tokens_in + response.tokens_out,
        }
        # Include cache tokens in usage if present (Anthropic)
        if response.cache_read_tokens or response.cache_creation_tokens:
            usage_block["prompt_tokens_details"] = {
                "cached_tokens": response.cache_read_tokens,
                "cache_creation_tokens": response.cache_creation_tokens,
            }

        openai_response = response.raw if response.raw else {
            "id": f"ecate-{int(time.time())}",
            "object": "chat.completion",
            "created": int(time.time()),
            "model": response.model,
            "choices": [{
                "index": 0,
                "message": {
                    "role": "assistant",
                    "content": response.content,
                    **({"tool_calls": response.tool_calls} if response.tool_calls else {}),
                },
                "finish_reason": response.finish_reason or "stop",
            }],
            "usage": usage_block,
        }

        # Calculate cost including cache tokens
        cost = _calculate_cost(
            decision.model,
            response.tokens_in,
            response.tokens_out,
            response.cache_read_tokens,
            response.cache_creation_tokens,
        )

        self.usage_logger.log(UsageRecord(
            api_key_id=auth.key_id,
            model_id=model_id,
            strategy=decision.strategy,
            tokens_in=response.tokens_in,
            tokens_out=response.tokens_out,
            cache_read_tokens=response.cache_read_tokens,
            cache_creation_tokens=response.cache_creation_tokens,
            cost_usd=cost,
            latency_ttfb_ms=elapsed,
            latency_total_ms=elapsed,
            was_fallback=was_fallback,
            fallback_from=fallback_from,
        ))

        headers = {"X-Ecate-Routed-To": model_id, "X-Ecate-Strategy": decision.strategy}
        if was_fallback:
            headers["X-Ecate-Fallback-From"] = fallback_from or ""
        return web.json_response(openai_response, headers=headers)

    async def _handle_stream(
        self, request: web.Request, bridge, bridge_request, decision, auth, start_time, model_id,
        was_fallback: bool = False, fallback_from: str | None = None,
    ) -> web.StreamResponse:
        headers = {
            "Content-Type": "text/event-stream",
            "Cache-Control": "no-cache",
            "X-Ecate-Routed-To": model_id,
            "X-Ecate-Strategy": decision.strategy,
        }
        if was_fallback:
            headers["X-Ecate-Fallback-From"] = fallback_from or ""

        resp = web.StreamResponse(status=200, headers=headers)
        await resp.prepare(request)

        ttfb = None
        stream_result = await bridge.stream(bridge_request, decision.model)
        
        try:
            async for line in stream_result:
                if ttfb is None:
                    ttfb = (time.monotonic() - start_time) * 1000
                await resp.write(f"{line}\n\n".encode())

            self.circuit_breaker.record_success(model_id)
        except Exception as e:
            self.circuit_breaker.record_failure(model_id)
            logger.error(f"Stream error for {model_id}: {e}")
            raise

        elapsed = (time.monotonic() - start_time) * 1000

        # Get usage from stream result (populated during iteration)
        usage = stream_result.usage

        # Calculate cost including cache tokens
        cost = _calculate_cost(
            decision.model,
            usage.tokens_in,
            usage.tokens_out,
            usage.cache_read_tokens,
            usage.cache_creation_tokens,
        )

        self.usage_logger.log(UsageRecord(
            api_key_id=auth.key_id,
            model_id=model_id,
            strategy=decision.strategy,
            tokens_in=usage.tokens_in,
            tokens_out=usage.tokens_out,
            cache_read_tokens=usage.cache_read_tokens,
            cache_creation_tokens=usage.cache_creation_tokens,
            cost_usd=cost,
            latency_ttfb_ms=ttfb,
            latency_total_ms=elapsed,
            was_fallback=was_fallback,
            fallback_from=fallback_from,
        ))

        return resp

    async def handle_models(self, request: web.Request) -> web.Response:
        """Handle GET /v1/models — list available models."""
        models = []
        for provider, model in self.config.all_models():
            models.append({
                "id": f"{provider.id}/{model.id}",
                "object": "model",
                "owned_by": provider.id,
                "permission": [],
            })
        # Add the auto routing models
        for strategy in ("auto", "auto:best", "auto:cheap"):
            models.append({
                "id": strategy,
                "object": "model",
                "owned_by": "ecate",
                "permission": [],
            })

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