"""Handler for chat completions endpoint."""

from __future__ import annotations

import asyncio
import time
import uuid
from collections.abc import AsyncGenerator

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

from robot_resources.keys import get_provider_key
from robot_resources.proxy.models import (
    ChatCompletionChunk,
    ChatCompletionRequest,
    ChatCompletionResponse,
    Choice,
    Message,
    RoutingInfo,
    StreamChoice,
    StreamDelta,
    Usage,
)
from robot_resources.proxy.providers import (
    AnthropicProvider,
    BaseProvider,
    GoogleProvider,
    OpenAIProvider,
    ProviderConfig,
)
from robot_resources.proxy.tokenizer import count_messages_tokens, count_tokens
from robot_resources.routing import MODELS_DB, async_route_prompt
from robot_resources.routing.decision_log import DecisionLog, RoutingDecision
from robot_resources.tracking.calculator import CostCalculator
from robot_resources.tracking.db import OutcomeDB
from robot_resources.tracking.recorder import OutcomeRecorder
from robot_resources.tracking.telemetry import get_telemetry_reporter

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

# ------------------------------------------------------------------
# Outcome tracking (lazy-initialized singleton)
# ------------------------------------------------------------------

_recorder: OutcomeRecorder | None = None
_outcome_db: OutcomeDB | None = None
_decision_log: DecisionLog | None = None
_baseline_cache: tuple[str | None, float] | None = None
_BASELINE_CACHE_TTL = 300.0  # 5 minutes


def _get_decision_log() -> DecisionLog | None:
    """Get the decision log, initializing on first call.

    Returns None if initialization fails — decision logging is optional
    and must never block request handling.
    """
    global _decision_log  # noqa: PLW0603
    if _decision_log is not None:
        return _decision_log
    try:
        _decision_log = DecisionLog()
        _decision_log.connect()
        logger.info("decision_log_initialized")
        return _decision_log
    except Exception:  # Broad catch intentional: fire-and-forget init (SC-01 exempt)
        logger.warning("decision_log_init_failed", exc_info=True)
        return None


def get_recorder() -> OutcomeRecorder | None:
    """Get the outcome recorder, initializing on first call.

    Returns None if initialization fails — tracking is optional
    and must never block request handling.
    """
    global _recorder, _outcome_db  # noqa: PLW0603
    if _recorder is not None:
        return _recorder
    try:
        _outcome_db = OutcomeDB()
        _outcome_db.connect()
        calculator = CostCalculator()
        _recorder = OutcomeRecorder(db=_outcome_db, calculator=calculator)
        logger.info("outcome_recorder_initialized")
        return _recorder
    except Exception:  # Broad catch intentional: fire-and-forget init (SC-01 exempt)
        logger.warning("outcome_recorder_init_failed", exc_info=True)
        return None


def _get_personal_baseline() -> dict | None:
    """Look up the user's most common direct model as personal baseline.

    Queries outcomes.db for the most frequently used non-routed model.
    Cached for 5 minutes to avoid querying on every request.
    Returns the model dict from MODELS_DB, or None if unavailable.
    """
    global _baseline_cache  # noqa: PLW0603
    now = time.monotonic()
    if _baseline_cache is not None and (now - _baseline_cache[1]) < _BASELINE_CACHE_TTL:
        model_name = _baseline_cache[0]
    else:
        # Need the outcome DB — reuse the recorder's DB if available
        recorder = get_recorder()
        if recorder is None or _outcome_db is None:
            return None
        valid = {m["name"] for m in MODELS_DB}
        model_name = _outcome_db.get_user_baseline_model(valid_models=valid)
        _baseline_cache = (model_name, now)

    if model_name is None:
        return None
    return next((m for m in MODELS_DB if m["name"] == model_name), None)


def _extract_request_keys(request: Request) -> dict[str, str]:
    """Extract provider API keys from incoming request headers.

    Supports transparent-proxy mode: the caller's own API key is read
    from the request and forwarded to the upstream provider.

    Skips the ``Authorization: Bearer`` token if it matches the
    Router's own ``ROUTER_API_KEY`` (used for Router auth, not a
    provider key).

    Returns a dict mapping provider name to API key.
    """
    from robot_resources.config import get_settings

    keys: dict[str, str] = {}
    router_api_key = get_settings().router_api_key or ""

    # Anthropic: x-api-key header
    anthropic_key = request.headers.get("x-api-key")
    if anthropic_key:
        keys["anthropic"] = anthropic_key

    # OpenAI: Authorization: Bearer <key>
    # Skip if the Bearer token is the Router's own auth key
    auth_header = request.headers.get("authorization", "")
    if auth_header.startswith("Bearer "):
        bearer_key = auth_header[len("Bearer "):]
        if bearer_key and bearer_key != router_api_key:
            keys["openai"] = bearer_key

    # Google: x-goog-api-key header
    google_key = request.headers.get("x-goog-api-key")
    if google_key:
        keys["google"] = google_key

    return keys


def _get_available_providers(request_keys: dict[str, str] | None = None) -> set[str]:
    """Get set of provider names that have API keys configured.

    Considers both stored keys (env/config) and per-request header keys.
    """
    available = set()
    for provider_name, (_, env_var) in PROVIDER_REGISTRY.items():
        if get_provider_key(env_var):
            available.add(provider_name)
    if request_keys:
        available.update(request_keys.keys())
    return available


def _get_available_models(request_keys: dict[str, str] | None = None) -> list[dict]:
    """Get models filtered to only providers with API keys."""
    available_providers = _get_available_providers(request_keys)
    if not available_providers:
        return MODELS_DB  # Return all if none available (will use placeholder)
    return [m for m in MODELS_DB if m["provider"] in available_providers]


# Provider registry: maps provider name to (ProviderClass, env_var)
PROVIDER_REGISTRY: dict[str, tuple[type[BaseProvider], str]] = {
    "anthropic": (AnthropicProvider, "ANTHROPIC_API_KEY"),
    "openai": (OpenAIProvider, "OPENAI_API_KEY"),
    "google": (GoogleProvider, "GOOGLE_API_KEY"),
}

# Cache for provider instances
_provider_cache: dict[str, BaseProvider] = {}


def get_provider(provider_name: str) -> BaseProvider | None:
    """
    Get a provider instance by name (using stored keys).

    Returns None if the API key is not configured.
    """
    if provider_name in _provider_cache:
        return _provider_cache[provider_name]

    if provider_name not in PROVIDER_REGISTRY:
        logger.warning("unknown_provider", provider=provider_name)
        return None

    provider_class, env_var = PROVIDER_REGISTRY[provider_name]
    api_key = get_provider_key(env_var)

    if not api_key:
        return None

    config = ProviderConfig(api_key=api_key)
    provider = provider_class(config)
    _provider_cache[provider_name] = provider

    logger.info("provider_initialized", provider=provider_name)
    return provider


def _get_provider_with_request_keys(
    provider_name: str,
    request_keys: dict[str, str],
) -> tuple[BaseProvider | None, str | None]:
    """Get a provider instance, preferring per-request header keys.

    Returns ``(provider, api_key_override)`` where *api_key_override*
    is the request-header key (passed to provider.complete/stream)
    or ``None`` when using a stored key.

    Per-request keys create ephemeral providers that are NOT cached.
    """
    if provider_name not in PROVIDER_REGISTRY:
        return None, None

    request_key = request_keys.get(provider_name)
    if request_key:
        # Ephemeral provider — do NOT cache (per-request key)
        provider_class, _ = PROVIDER_REGISTRY[provider_name]
        provider = provider_class(ProviderConfig(api_key=request_key))
        return provider, request_key

    # Fall back to stored key
    return get_provider(provider_name), None


def _infer_provider(model: str) -> str | None:
    """Infer provider from model name when not routing."""
    model_lower = model.lower()
    if "gpt" in model_lower or model_lower.startswith(("o1", "o3", "o4")):
        return "openai"
    if "claude" in model_lower:
        return "anthropic"
    if "gemini" in model_lower:
        return "google"
    return None


def _extract_prompt(request: ChatCompletionRequest) -> str:
    """Extract the prompt text from request messages for routing."""
    if not request.messages:
        return ""

    # Concatenate user messages for task detection
    user_messages = [m.content for m in request.messages if m.role == "user"]
    if user_messages:
        return user_messages[-1]  # Use last user message for routing

    # Fallback to last message
    return request.messages[-1].content


def _should_route(request: ChatCompletionRequest) -> bool:
    """Determine if request should be routed or use requested model directly."""
    # Route if model is "auto" or starts with "rr-" (robot resources prefix)
    model = request.model.lower()
    return model == "auto" or model.startswith("rr-")


@router.post("/v1/chat/completions", response_model=None)
async def create_chat_completion(
    request: ChatCompletionRequest,
    raw_request: Request,
) -> ChatCompletionResponse | StreamingResponse:
    """
    Create a chat completion with intelligent routing.

    If model is "auto" or starts with "rr-", Robot Resources will
    automatically select the optimal model based on task detection.
    Otherwise, the requested model is used directly.

    Supports transparent-proxy mode: API keys are extracted from
    the incoming request headers and forwarded to upstream providers.
    """
    # Extract per-request API keys from headers (transparent proxy)
    request_keys = _extract_request_keys(raw_request)

    # Extract prompt and route if needed
    prompt = _extract_prompt(request)
    routing_decision = None
    routed_model = request.model

    if _should_route(request):
        # Only route to models from providers with available API keys
        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"]

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

    logger.info(
        "chat_completion_request",
        model=routed_model,
        original_model=request.model,
        messages_count=len(request.messages),
        stream=request.stream,
        routed=routing_decision is not None,
        passthrough_keys=list(request_keys.keys()) if request_keys else [],
    )

    # Build routing info if we routed
    routing_info = None
    if routing_decision:
        routing_info = RoutingInfo(
            selected_model=routing_decision["selected_model"],
            original_model=request.model,
            provider=routing_decision["provider"],
            task_type=routing_decision["task_type"],
            capability_score=routing_decision["capability_score"],
            savings_percent=routing_decision["savings_percent"],
            baseline_model=routing_decision["baseline_model"],
            reasoning=routing_decision["reasoning"],
        )

    # Get provider for the routed model (request-header keys take priority)
    if routing_decision:
        provider_name = routing_decision["provider"]
    else:
        provider_name = _infer_provider(routed_model)

    api_key_override: str | None = None
    if provider_name:
        provider, api_key_override = _get_provider_with_request_keys(
            provider_name, request_keys
        )
    else:
        provider = None

    if request.stream:
        return StreamingResponse(
            _stream_response(
                request, routed_model, routing_decision, provider, provider_name,
                prompt=prompt, api_key_override=api_key_override,
                request_keys=request_keys,
            ),
            media_type="text/event-stream",
        )

    # Non-streaming response — try with retry + cross-provider fallback
    t_start = time.monotonic()

    result = await _call_with_retry_and_fallback(
        request=request,
        routed_model=routed_model,
        provider_name=provider_name,
        provider=provider,
        routing_decision=routing_decision,
        api_key_override=api_key_override,
        request_keys=request_keys,
    )

    if result is not None:
        response, actual_provider, actual_model = result
        response.routing_info = routing_info
        latency_ms = int((time.monotonic() - t_start) * 1000)

        logger.info(
            "chat_completion_forwarded",
            provider=actual_provider,
            model=response.model,
            tokens=response.usage.total_tokens,
            latency_ms=latency_ms,
            fallback=actual_model != routed_model,
        )

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

        return response

    # All providers failed — placeholder response
    response_content = _generate_placeholder_response(request, routed_model, routing_decision)

    prompt_tok = _estimate_tokens(request.messages, routed_model)
    completion_tok = _estimate_tokens_str(response_content, routed_model)
    usage = Usage(
        prompt_tokens=prompt_tok,
        completion_tokens=completion_tok,
        total_tokens=prompt_tok + completion_tok,
    )
    response = ChatCompletionResponse(
        model=routed_model,
        choices=[
            Choice(
                index=0,
                message=Message(role="assistant", content=response_content),
                finish_reason="stop",
            )
        ],
        usage=usage,
        routing_info=routing_info,
    )
    latency_ms = int((time.monotonic() - t_start) * 1000)

    logger.info(
        "chat_completion_placeholder",
        model=response.model,
        tokens=response.usage.total_tokens,
        reason="all_providers_failed",
    )

    _schedule_tracking(
        routing_decision=routing_decision,
        provider_name=provider_name,
        model=routed_model,
        usage=usage,
        latency_ms=latency_ms,
        success=None,  # placeholder — unknown success
        prompt=prompt,
    )

    return response


# ------------------------------------------------------------------
# Retry + cross-provider fallback
# ------------------------------------------------------------------

_MAX_RETRIES = 2  # 1 initial + 1 retry
_BACKOFF_BASE = 0.5  # seconds (0.5s, 1.0s)


async def _attempt_provider_call(
    request: ChatCompletionRequest,
    provider: BaseProvider,
    model: str,
    *,
    api_key: str | None = None,
) -> ChatCompletionResponse:
    """Attempt a single provider call.  Raises on failure."""
    forwarded = ChatCompletionRequest(
        model=model,
        messages=request.messages,
        temperature=request.temperature,
        max_tokens=request.max_tokens,
        stream=False,
    )
    return await provider.complete(forwarded, api_key=api_key)


async def _call_with_retry_and_fallback(
    *,
    request: ChatCompletionRequest,
    routed_model: str,
    provider_name: str | None,
    provider: BaseProvider | None,
    routing_decision: dict | None,
    api_key_override: str | None = None,
    request_keys: dict[str, str] | None = None,
) -> tuple[ChatCompletionResponse, str, str] | None:
    """Try primary provider with retry, then fall back to alternatives.

    Returns ``(response, actual_provider, actual_model)`` on success,
    or ``None`` if every attempt fails.
    """
    # --- Primary provider with retry ---
    if provider:
        for attempt in range(_MAX_RETRIES):
            try:
                response = await _attempt_provider_call(
                    request, provider, routed_model, api_key=api_key_override,
                )
                return (response, provider_name or "unknown", routed_model)
            except (httpx.HTTPStatusError, httpx.RequestError, httpx.TimeoutException) as e:
                logger.warning(
                    "provider_attempt_failed",
                    provider=provider_name,
                    model=routed_model,
                    attempt=attempt + 1,
                    max_retries=_MAX_RETRIES,
                    error=str(e),
                )
                if attempt < _MAX_RETRIES - 1:
                    await asyncio.sleep(_BACKOFF_BASE * (2**attempt))

    # --- Fallback to next-cheapest model from a DIFFERENT provider ---
    candidates = (routing_decision or {}).get("ranked_candidates", [])
    failed_provider = provider_name

    for candidate in candidates:
        if candidate["provider"] == failed_provider:
            continue  # Same provider likely has the same outage

        fb_key: str | None = None
        if request_keys:
            fallback_provider, fb_key = _get_provider_with_request_keys(
                candidate["provider"], request_keys
            )
        else:
            fallback_provider = get_provider(candidate["provider"])
        if fallback_provider is None:
            continue  # No API key configured

        try:
            response = await _attempt_provider_call(
                request, fallback_provider, candidate["name"], api_key=fb_key,
            )
            logger.info(
                "fallback_provider_succeeded",
                original_provider=failed_provider,
                fallback_provider=candidate["provider"],
                fallback_model=candidate["name"],
            )
            return (response, candidate["provider"], candidate["name"])
        except (httpx.HTTPStatusError, httpx.RequestError, httpx.TimeoutException) as e:
            logger.warning(
                "fallback_attempt_failed",
                provider=candidate["provider"],
                model=candidate["name"],
                error=str(e),
            )

    return None


async def _stream_response(
    request: ChatCompletionRequest,
    routed_model: str,
    routing_decision: dict | None,
    provider: BaseProvider | None = None,
    provider_name: str | None = None,
    prompt: str = "",
    api_key_override: str | None = None,
    request_keys: dict[str, str] | None = None,
) -> AsyncGenerator[str, None]:
    """Generate streaming response chunks.

    If a provider is available, streams real responses from the LLM.
    Falls back to placeholder streaming when no provider or on error.
    """
    t_start = time.monotonic()

    # Build list of providers to try: primary first, then fallback candidates
    # Each entry: (provider, provider_name, model, api_key_override)
    providers_to_try: list[tuple[BaseProvider, str, str, str | None]] = []
    if provider and provider_name:
        providers_to_try.append((provider, provider_name, routed_model, api_key_override))

    # Add cross-provider fallback candidates
    candidates = (routing_decision or {}).get("ranked_candidates", [])
    for candidate in candidates:
        if candidate["provider"] == provider_name:
            continue
        fb_key: str | None = None
        if request_keys:
            fb_provider, fb_key = _get_provider_with_request_keys(
                candidate["provider"], request_keys
            )
        else:
            fb_provider = get_provider(candidate["provider"])
        if fb_provider is not None:
            providers_to_try.append((fb_provider, candidate["provider"], candidate["name"], fb_key))

    for try_provider, try_provider_name, try_model, try_api_key in providers_to_try:
        try:
            forwarded_request = ChatCompletionRequest(
                model=try_model,
                messages=request.messages,
                temperature=request.temperature,
                max_tokens=request.max_tokens,
                stream=True,
            )

            async for chunk_str in try_provider.stream(forwarded_request, api_key=try_api_key):
                yield chunk_str

            latency_ms = int((time.monotonic() - t_start) * 1000)

            usage = try_provider._last_stream_usage
            if not usage:
                usage = Usage(
                    prompt_tokens=_estimate_tokens(request.messages, try_model),
                    completion_tokens=0,
                    total_tokens=_estimate_tokens(request.messages, try_model),
                )

            logger.info(
                "chat_completion_streamed",
                provider=try_provider_name,
                model=try_model,
                tokens=usage.total_tokens,
                latency_ms=latency_ms,
                fallback=try_model != routed_model,
            )

            _schedule_tracking(
                routing_decision=routing_decision,
                provider_name=try_provider_name,
                model=try_model,
                usage=usage,
                latency_ms=latency_ms,
                success=True,
                prompt=prompt,
            )
            return

        except (httpx.HTTPStatusError, httpx.RequestError, httpx.TimeoutException) as e:
            logger.warning(
                "provider_stream_error",
                provider=try_provider_name,
                model=try_model,
                error=str(e),
            )
            # Try next provider

    # Fallback: placeholder streaming (all providers failed)
    chunk_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    created = int(time.time())

    first_chunk = ChatCompletionChunk(
        id=chunk_id,
        created=created,
        model=routed_model,
        choices=[
            StreamChoice(
                index=0,
                delta=StreamDelta(role="assistant", content=""),
                finish_reason=None,
            )
        ],
    )
    yield f"data: {first_chunk.model_dump_json()}\n\n"

    content = _generate_placeholder_response(request, routed_model, routing_decision)
    words = content.split()

    for i, word in enumerate(words):
        chunk_content = word + (" " if i < len(words) - 1 else "")
        chunk = ChatCompletionChunk(
            id=chunk_id,
            created=created,
            model=routed_model,
            choices=[
                StreamChoice(
                    index=0,
                    delta=StreamDelta(content=chunk_content),
                    finish_reason=None,
                )
            ],
        )
        yield f"data: {chunk.model_dump_json()}\n\n"

    final_chunk = ChatCompletionChunk(
        id=chunk_id,
        created=created,
        model=routed_model,
        choices=[
            StreamChoice(
                index=0,
                delta=StreamDelta(),
                finish_reason="stop",
            )
        ],
    )
    yield f"data: {final_chunk.model_dump_json()}\n\n"
    yield "data: [DONE]\n\n"


def _generate_placeholder_response(
    request: ChatCompletionRequest,
    routed_model: str,
    routing_decision: dict | None,
) -> str:
    """Generate a placeholder response showing routing decision."""
    last_message = request.messages[-1].content if request.messages else ""

    if routing_decision:
        return (
            f"[Robot Resources] Routed to {routed_model} "
            f"(task: {routing_decision['task_type']}, "
            f"savings: {routing_decision['savings_percent']}%). "
            f"Prompt: '{last_message[:30]}...'"
        )

    return f"[Robot Resources] Using {routed_model}. Prompt: '{last_message[:50]}...'"


def _estimate_tokens(messages: list[Message], model: str = "") -> int:
    """Estimate tokens for a list of messages using tiktoken."""
    return count_messages_tokens(messages, model)


def _estimate_tokens_str(text: str, model: str = "") -> int:
    """Estimate tokens for a string using tiktoken."""
    return count_tokens(text, model)


# ------------------------------------------------------------------
# Tracking integration (non-blocking, fire-and-forget)
# ------------------------------------------------------------------


def _schedule_tracking(
    *,
    routing_decision: dict | None,
    provider_name: str | None,
    model: str,
    usage: Usage,
    latency_ms: int,
    success: bool | None,
    prompt: str = "",
) -> None:
    """Schedule outcome recording and decision logging as background tasks.

    Never blocks the response. If the recorder is unavailable or
    the task fails, errors are logged and swallowed.
    """
    recorder = get_recorder()
    if recorder is None:
        return

    try:
        if routing_decision is not None:
            record = recorder.build_record(
                routing_decision=routing_decision,
                input_tokens=usage.prompt_tokens,
                output_tokens=usage.completion_tokens,
                latency_ms=latency_ms,
                success=success,
            )
        else:
            record = recorder.build_record_direct(
                model=model,
                provider=provider_name or "unknown",
                input_tokens=usage.prompt_tokens,
                output_tokens=usage.completion_tokens,
                latency_ms=latency_ms,
                success=success,
            )
        asyncio.create_task(recorder.arecord_safe(record))

        # Decision logging (fire-and-forget alongside outcome recording)
        if routing_decision is not None:
            _schedule_decision_log(routing_decision, prompt, latency_ms)

        # Platform telemetry (fire-and-forget alongside local recording)
        telemetry = get_telemetry_reporter()
        if telemetry.enabled:
            event = telemetry.build_event(
                selected_model=record.model_used,
                provider=record.provider,
                task_type=record.task_type,
                cost_actual=record.cost_actual,
                cost_baseline=record.cost_baseline,
                savings_percent=record.savings_percent or 0.0,
                cost_saved=record.cost_saved,
                latency_ms=record.latency_ms,
                input_tokens=record.input_tokens,
                output_tokens=record.output_tokens,
                savings_mode="cost",
            )
            asyncio.create_task(telemetry.areport_safe(event))
    except Exception:  # Broad catch intentional: fire-and-forget tracking (SC-01 exempt)
        logger.warning("tracking_schedule_failed", exc_info=True)


def _schedule_decision_log(
    routing_decision: dict, prompt: str, latency_ms: int
) -> None:
    """Log a routing decision to SQLite. Fire-and-forget."""
    dl = _get_decision_log()
    if dl is None:
        return
    try:
        import hashlib

        decision = RoutingDecision(
            prompt_hash=hashlib.sha256(prompt.encode()).hexdigest()[:16],
            prompt_preview=prompt[:100],
            keyword_task_type=routing_decision.get(
                "keyword_task_type", routing_decision["task_type"]
            ),
            keyword_confidence=routing_decision.get(
                "keyword_confidence",
                routing_decision.get("detection_confidence", 0.0),
            ),
            classification_source=routing_decision.get("classification_source", "keyword"),
            llm_task_type=routing_decision.get("llm_task_type"),
            llm_complexity=routing_decision.get("llm_complexity"),
            llm_model=routing_decision.get("llm_model"),
            final_task_type=routing_decision["task_type"],
            final_model=routing_decision["selected_model"],
            capability_threshold=routing_decision.get("capability_threshold", 0.7),
            latency_ms=latency_ms,
        )
        asyncio.create_task(dl.alog(decision))
    except Exception:  # Broad catch intentional: fire-and-forget (SC-01 exempt)
        logger.warning("decision_log_schedule_failed", exc_info=True)
