"""
kvict Proxy — FastAPI Application.

Drop-in replacement for OpenAI's /v1/chat/completions endpoint.
Routes requests through CAAE-inspired caching, tracks token savings,
and exposes metrics for the dashboard and browser extension.

Usage:
    uvicorn proxy.app:app --host 0.0.0.0 --port 8080

Architecture:
    Request -> Cache check -> (hit: return cached) / (miss: forward to provider)
              -> Record metrics -> Return response

    CAAE integration points:
    1. Cache eviction uses cost-model scoring (expensive responses survive longer)
    2. Circuit breaker serves stale cache when provider is degraded
"""

from __future__ import annotations

import json
import logging
import time
from contextlib import asynccontextmanager
from typing import Any, Dict

import httpx
from fastapi import FastAPI, Request, Response
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, StreamingResponse
from pydantic import ValidationError

from proxy.adapters import get_adapter
from proxy.adapters.base import ProviderAdapter
from proxy.cache import ResponseCache
from proxy.circuit_breaker import ProviderCircuitBreaker
from proxy.config import ProxyConfig, get_config
from proxy.feature_flags import get_feature_flags
from proxy.savings_api import router as savings_router, set_store
from proxy.schemas import ChatCompletionRequest, ChatCompletionResponse
from proxy.store import MetricsStore
from proxy.streaming import stream_provider_response, stream_result_to_response
from proxy.token_tracker import compute_savings, count_tokens

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Globals (initialized in lifespan)
# ---------------------------------------------------------------------------
_adapter: ProviderAdapter | None = None
_cache: ResponseCache | None = None
_store: MetricsStore | None = None
_circuit_breaker: ProviderCircuitBreaker | None = None
_http_client: httpx.AsyncClient | None = None


@asynccontextmanager
async def lifespan(application: FastAPI):
    """Initialize shared resources on startup, clean up on shutdown."""
    global _adapter, _cache, _store, _circuit_breaker, _http_client

    config = get_config()
    flags = get_feature_flags()
    _adapter = get_adapter(config.provider, config.provider_base_url)
    _cache = ResponseCache(config)
    _store = MetricsStore(config)
    _circuit_breaker = (
        ProviderCircuitBreaker(config) if flags.circuit_breaker_enabled else None
    )
    _http_client = httpx.AsyncClient(timeout=httpx.Timeout(60.0, connect=10.0))

    # Inject store into savings API router
    set_store(_store)

    logger.info(
        "kvict Proxy started: provider=%s, cache_ttl=%ds, circuit_breaker_threshold=%dms",
        _adapter.name,
        config.cache_ttl_seconds,
        config.circuit_breaker_threshold_ms,
    )

    yield

    # Cleanup
    if _http_client:
        await _http_client.aclose()
    logger.info("kvict Proxy shut down")


# ---------------------------------------------------------------------------
# App
# ---------------------------------------------------------------------------
app = FastAPI(
    title="kvict Proxy",
    description=(
        "Subscription stretching proxy for LLM APIs. "
        "Drop-in replacement for /v1/chat/completions with CAAE-powered caching."
    ),
    version="1.0.0",
    lifespan=lifespan,
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Mount savings/metrics API
app.include_router(savings_router)


# ---------------------------------------------------------------------------
# Health
# ---------------------------------------------------------------------------
@app.get("/health")
async def health():
    return {
        "status": "healthy",
        "provider": _adapter.name if _adapter else "uninitialized",
        "cache_size": _cache.size if _cache else 0,
        "circuit_breaker": (
            _circuit_breaker.state.value if _circuit_breaker else "unknown"
        ),
    }


# ---------------------------------------------------------------------------
# Main proxy endpoint
# ---------------------------------------------------------------------------
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
    """Drop-in replacement for OpenAI's /v1/chat/completions.

    Flow:
    1. Parse request
    2. Check cache (exact match, CAAE-scored eviction)
    3. Check circuit breaker (serve stale if provider degraded)
    4. Forward to provider via adapter
    5. Record metrics (token savings)
    6. Return response
    """
    assert _adapter is not None
    assert _cache is not None
    assert _store is not None
    assert _http_client is not None

    # Extract API key from Authorization header (passthrough, never stored)
    auth_header = request.headers.get("authorization", "")
    api_key = ""
    if auth_header.startswith("Bearer "):
        api_key = auth_header[7:]
    elif auth_header.startswith("x-api-key "):
        api_key = auth_header[10:]

    # Parse request body (return 422 on validation error for missing/invalid fields)
    body = await request.json()
    try:
        chat_request = ChatCompletionRequest(**body)
    except ValidationError as e:
        return JSONResponse(
            status_code=422,
            content={"detail": e.errors()},
        )
    is_streaming = chat_request.stream or False

    # Count baseline tokens (what user WOULD pay without proxy)
    baseline_tokens = count_tokens(chat_request.messages, chat_request.model)

    # Get pricing for this model
    pricing = _adapter.get_pricing(chat_request.model)

    # ---- Step 1: Check cache ----
    if not is_streaming:
        cached_response = _cache.get(chat_request)
        if cached_response is not None:
            return _handle_cache_hit(
                cached_response, chat_request, baseline_tokens, pricing
            )

    # ---- Step 2: Check circuit breaker ----
    if not _circuit_breaker.should_allow_request():
        # Provider degraded — try to serve stale cache
        if not is_streaming:
            stale = _cache.get(chat_request)
            if stale is not None:
                _circuit_breaker.record_stale_served()
                logger.info("Serving stale cache (circuit breaker OPEN)")
                return _handle_cache_hit(stale, chat_request, baseline_tokens, pricing)
        # No cached response — must try provider anyway
        logger.warning(
            "Circuit breaker OPEN but no cached response available; forwarding to provider"
        )

    # ---- Step 3: Forward to provider ----
    url, headers, req_body = _adapter.translate_request(chat_request, api_key)
    start_time = time.time()

    try:
        if is_streaming:
            return await _handle_streaming(
                url,
                headers,
                req_body,
                chat_request,
                baseline_tokens,
                pricing,
                start_time,
            )
        else:
            return await _handle_non_streaming(
                url,
                headers,
                req_body,
                chat_request,
                baseline_tokens,
                pricing,
                start_time,
            )
    except (httpx.TimeoutException, httpx.ConnectError, httpx.HTTPStatusError) as exc:
        latency_ms = (time.time() - start_time) * 1000
        if _circuit_breaker:
            _circuit_breaker.record_failure()
        logger.error("Provider error (%.0fms): %s", latency_ms, exc)

        # Try stale cache as fallback
        if not is_streaming:
            stale = _cache.get(chat_request)
            if stale is not None:
                if _circuit_breaker:
                    _circuit_breaker.record_stale_served()
                logger.info("Serving stale cache after provider error")
                return _handle_cache_hit(stale, chat_request, baseline_tokens, pricing)

        return JSONResponse(
            status_code=502,
            content={
                "error": {"message": f"Provider error: {exc}", "type": "proxy_error"}
            },
        )


# ---------------------------------------------------------------------------
# Handlers
# ---------------------------------------------------------------------------


def _handle_cache_hit(
    response: ChatCompletionResponse,
    request: ChatCompletionRequest,
    baseline_tokens: int,
    pricing,
) -> JSONResponse:
    """Handle a cache hit: record savings and return cached response."""
    completion_tokens = response.usage.completion_tokens if response.usage else 0
    savings = compute_savings(
        baseline_tokens=baseline_tokens,
        actual_tokens=0,
        completion_tokens=completion_tokens,
        pricing=pricing,
        cache_hit=True,
    )

    baseline_cost = (
        baseline_tokens * pricing.input_cost_per_1k / 1000
        + completion_tokens * pricing.output_cost_per_1k / 1000
    )

    _store.record_request(
        provider=_adapter.name,
        model=request.model,
        prompt_tokens_actual=0,
        prompt_tokens_baseline=baseline_tokens,
        completion_tokens=completion_tokens,
        cached_tokens=savings.cached_tokens,
        cache_hit=True,
        latency_ms=0.0,
        cost_actual_usd=0.0,
        cost_baseline_usd=baseline_cost,
    )

    return JSONResponse(
        content=response.model_dump(),
        headers={"x-kvict-cache-hit": "true"},
    )


async def _handle_non_streaming(
    url: str,
    headers: Dict[str, str],
    req_body: Dict[str, Any],
    chat_request: ChatCompletionRequest,
    baseline_tokens: int,
    pricing,
    start_time: float,
) -> JSONResponse:
    """Forward non-streaming request to provider, cache result, record metrics."""
    provider_response = await _http_client.post(url, headers=headers, json=req_body)
    latency_ms = (time.time() - start_time) * 1000

    if provider_response.status_code != 200:
        _circuit_breaker.record_failure()
        return JSONResponse(
            status_code=provider_response.status_code,
            content=provider_response.json(),
        )

    # Success — record in circuit breaker
    if _circuit_breaker:
        _circuit_breaker.record_success(latency_ms)

    response_body = provider_response.json()
    response = _adapter.translate_response(
        provider_response.status_code,
        dict(provider_response.headers),
        response_body,
    )

    actual_tokens = response.usage.prompt_tokens if response.usage else baseline_tokens
    completion_tokens = response.usage.completion_tokens if response.usage else 0

    savings = compute_savings(
        baseline_tokens=baseline_tokens,
        actual_tokens=actual_tokens,
        completion_tokens=completion_tokens,
        pricing=pricing,
        cache_hit=False,
    )

    actual_cost = (
        actual_tokens * pricing.input_cost_per_1k / 1000
        + completion_tokens * pricing.output_cost_per_1k / 1000
    )
    baseline_cost = (
        baseline_tokens * pricing.input_cost_per_1k / 1000
        + completion_tokens * pricing.output_cost_per_1k / 1000
    )

    # Store in cache (CAAE cost-model-scored eviction applies here)
    _cache.put(
        request=chat_request,
        response=response,
        prompt_tokens=actual_tokens,
        completion_tokens=completion_tokens,
        model_input_cost_per_1k=pricing.input_cost_per_1k,
    )

    # Record metrics
    _store.record_request(
        provider=_adapter.name,
        model=chat_request.model,
        prompt_tokens_actual=actual_tokens,
        prompt_tokens_baseline=baseline_tokens,
        completion_tokens=completion_tokens,
        cached_tokens=savings.cached_tokens,
        cache_hit=False,
        latency_ms=latency_ms,
        cost_actual_usd=actual_cost,
        cost_baseline_usd=baseline_cost,
    )

    return JSONResponse(content=response.model_dump())


async def _handle_streaming(
    url: str,
    headers: Dict[str, str],
    req_body: Dict[str, Any],
    chat_request: ChatCompletionRequest,
    baseline_tokens: int,
    pricing,
    start_time: float,
) -> StreamingResponse:
    """Forward streaming request, buffer for caching, record metrics on completion."""

    async def generate():
        try:
            async with _http_client.stream(
                "POST", url, headers=headers, json=req_body
            ) as provider_response:
                if provider_response.status_code != 200:
                    _circuit_breaker.record_failure()
                    error_body = await provider_response.aread()
                    yield f"data: {error_body.decode()}\n\n"
                    return

                latency_ms = (time.time() - start_time) * 1000
                _circuit_breaker.record_success(latency_ms)

                stream_result = None
                async for sse_line, result in stream_provider_response(
                    provider_response, _adapter
                ):
                    stream_result = result
                    yield sse_line

                # Stream complete — cache the full response and record metrics
                if stream_result is not None:
                    full_response = stream_result_to_response(stream_result)
                    actual_tokens = stream_result.prompt_tokens or baseline_tokens
                    completion_tokens = stream_result.completion_tokens

                    _cache.put(
                        request=chat_request,
                        response=full_response,
                        prompt_tokens=actual_tokens,
                        completion_tokens=completion_tokens,
                        model_input_cost_per_1k=pricing.input_cost_per_1k,
                    )

                    savings = compute_savings(
                        baseline_tokens=baseline_tokens,
                        actual_tokens=actual_tokens,
                        completion_tokens=completion_tokens,
                        pricing=pricing,
                        cache_hit=False,
                    )

                    actual_cost = (
                        actual_tokens * pricing.input_cost_per_1k / 1000
                        + completion_tokens * pricing.output_cost_per_1k / 1000
                    )
                    baseline_cost = (
                        baseline_tokens * pricing.input_cost_per_1k / 1000
                        + completion_tokens * pricing.output_cost_per_1k / 1000
                    )

                    _store.record_request(
                        provider=_adapter.name,
                        model=chat_request.model,
                        prompt_tokens_actual=actual_tokens,
                        prompt_tokens_baseline=baseline_tokens,
                        completion_tokens=completion_tokens,
                        cached_tokens=savings.cached_tokens,
                        cache_hit=False,
                        latency_ms=(time.time() - start_time) * 1000,
                        cost_actual_usd=actual_cost,
                        cost_baseline_usd=baseline_cost,
                    )

        except (httpx.TimeoutException, httpx.ConnectError) as exc:
            if _circuit_breaker:
                _circuit_breaker.record_failure()
            logger.error("Streaming provider error: %s", exc)
            yield f'data: {{"error": {{"message": "Provider error: {exc}"}}}}\n\n'

    return StreamingResponse(generate(), media_type="text/event-stream")
