"""AI Proxy — FastAPI server exposing an OpenAI-compatible API.

The server wraps LiteLLM with:
  - Content-based routing (analyzer → router → cheapest adequate model)
  - Delta context buffer (inject code2llm diffs into system prompt)
  - Semantic caching (Redis/in-memory)
  - Cost tracking with budget enforcement
  - Fallback chains across providers

Start:
  proxym                    # uses .env defaults
  uvicorn proxym.main:app   # manual
"""

from __future__ import annotations

import argparse
import json
import time
from contextlib import asynccontextmanager

import litellm
import structlog
import uvicorn
from fastapi import FastAPI, Header, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse

from proxym.cache import DeltaBuffer
from proxym.config import get_settings
from proxym.middleware import AuthMiddleware, CostTrackingMiddleware
from proxym.providers import MODELS, get_model
from proxym.router import AnalysisResult, analyze_request
from proxym.router.strategy import CostLedger, Router, RoutingDecision

structlog.configure(
    processors=[
        structlog.stdlib.add_log_level,
        structlog.dev.ConsoleRenderer(),
    ],
    wrapper_class=structlog.stdlib.BoundLogger,
    context_class=dict,
    logger_factory=structlog.PrintLoggerFactory(),
)

logger = structlog.get_logger()

# ── Global state ─────────────────────────────────────────────

_router: Router | None = None
_delta_buffer: DeltaBuffer | None = None
_context_payload: str = ""  # latest context from delta watcher


def _get_router() -> Router:
    global _router
    if _router is None:
        settings = get_settings()
        providers = set(settings.available_providers().keys())
        ledger = CostLedger(
            daily_limit=settings.daily_budget_usd,
            monthly_limit=settings.monthly_budget_usd,
        )
        _router = Router(
            available_providers=providers,
            ledger=ledger,
            max_request_cost=settings.max_request_cost_usd,
        )
    return _router


def _get_delta_buffer() -> DeltaBuffer:
    global _delta_buffer
    if _delta_buffer is None:
        settings = get_settings()
        _delta_buffer = DeltaBuffer(
            watch_dir=settings.watch_dir,
            max_context_tokens=settings.max_context_tokens,
            delta_threshold=settings.delta_threshold,
        )
    return _delta_buffer


# ── App lifecycle ────────────────────────────────────────────

@asynccontextmanager
async def lifespan(app: FastAPI):
    settings = get_settings()
    providers = settings.available_providers()
    logger.info(
        "proxy_starting",
        port=settings.port,
        providers=list(providers.keys()),
        models=len(MODELS),
        budget_daily=f"${settings.daily_budget_usd}",
        budget_monthly=f"${settings.monthly_budget_usd}",
    )

    # Configure LiteLLM API keys
    if settings.anthropic_api_key:
        litellm.anthropic_key = settings.anthropic_api_key
    if settings.openai_api_key:
        litellm.openai_key = settings.openai_api_key

    # Set LiteLLM to drop unsupported params silently
    litellm.drop_params = True
    litellm.set_verbose = False

    yield

    logger.info("proxy_stopped")


app = FastAPI(
    title="Proxym",
    description="Intelligent multi-provider LLM gateway",
    version="0.1.0",
    lifespan=lifespan,
)

# Mount dashboard
from proxym.dashboard import router as dashboard_router
app.include_router(dashboard_router)

# Add middleware
settings = get_settings()
app.add_middleware(CostTrackingMiddleware)
app.add_middleware(AuthMiddleware, master_key=settings.master_key)


# ── Health & info endpoints ──────────────────────────────────

@app.get("/health")
async def health():
    return {"status": "ok", "version": "0.1.0"}


@app.get("/v1/models")
async def list_models():
    """List available models (OpenAI-compatible)."""
    models_list = []
    providers = set(get_settings().available_providers().keys())
    for mid, m in MODELS.items():
        if m.provider in providers:
            models_list.append({
                "id": mid,
                "object": "model",
                "owned_by": m.provider,
                "display_name": m.display_name,
                "input_cost_per_1m": m.input_cost,
                "output_cost_per_1m": m.output_cost,
                "context_window": m.context_window,
                "capabilities": [c.value for c in m.capabilities],
                "tier_range": f"{m.min_tier.value}–{m.max_tier.value}",
            })
    return {"object": "list", "data": models_list}


@app.get("/v1/budget")
async def budget_status():
    """Return current budget usage."""
    router = _get_router()
    return {
        "daily_remaining_usd": round(router.ledger.daily_remaining, 4),
        "monthly_remaining_usd": round(router.ledger.monthly_remaining, 4),
        "daily_limit_usd": router.ledger.daily_limit,
        "monthly_limit_usd": router.ledger.monthly_limit,
    }


@app.get("/v1/context/stats")
async def context_stats():
    """Return delta buffer statistics."""
    buf = _get_delta_buffer()
    return buf.get_stats()


# ── Context delta endpoint (receives from watcher client) ────

@app.post("/v1/context/delta")
async def receive_delta(request: Request):
    """Receive context delta from the file watcher client."""
    global _context_payload
    body = await request.json()
    _context_payload = body.get("payload", "")
    return {
        "ok": True,
        "version": body.get("version"),
        "token_estimate": body.get("token_estimate"),
    }


# ── Main completion endpoint (OpenAI-compatible) ─────────────

@app.post("/v1/chat/completions")
async def chat_completions(
    request: Request,
    x_task_tier: str | None = Header(None, alias="X-Task-Tier"),
):
    """OpenAI-compatible chat completions with intelligent routing.

    Extra headers:
      X-Task-Tier: trivial|operational|standard|complex|deep
      X-Inject-Context: true  (inject latest delta context)
    """
    body = await request.json()
    messages: list[dict] = body.get("messages", [])
    requested_model: str | None = body.get("model")
    stream: bool = body.get("stream", False)

    inject_context = request.headers.get("X-Inject-Context", "").lower() == "true"

    # 1. Inject delta context if requested and available
    if inject_context and _context_payload:
        # Prepend context as system message
        context_msg = {
            "role": "system",
            "content": f"Current codebase context:\n{_context_payload}",
        }
        # Insert after existing system message or at start
        if messages and messages[0].get("role") == "system":
            messages[0]["content"] += f"\n\n{_context_payload}"
        else:
            messages.insert(0, context_msg)

    # 2. Analyze content
    analysis = analyze_request(
        messages,
        explicit_tier=x_task_tier,
        explicit_model=requested_model if requested_model and "/" in requested_model else None,
    )

    # 3. Route to model
    router = _get_router()

    # Check if the requested model is a known model ID or an alias
    explicit = None
    if requested_model:
        # Allow tier aliases: "cheap", "balanced", "premium", "free", "local"
        tier_aliases = {
            "cheap": "anthropic/haiku-4.5",
            "balanced": "anthropic/sonnet-4.6",
            "premium": "anthropic/opus-4.6",
            "free": "google/gemini-2.5-flash",
            "local": "ollama/qwen2.5-coder-3b",
        }
        explicit = tier_aliases.get(requested_model, requested_model)

    decision = router.route(analysis, explicit_model=explicit)

    # 4. Build LiteLLM kwargs
    litellm_kwargs = {
        "model": decision.model.litellm_model,
        "messages": messages,
        "stream": stream,
    }

    # Pass through supported OpenAI params
    for key in ("temperature", "max_tokens", "top_p", "stop", "tools",
                "tool_choice", "response_format", "seed"):
        if key in body:
            litellm_kwargs[key] = body[key]

    # Set API keys per provider
    _set_provider_keys(litellm_kwargs, decision.model)

    # 5. Call LiteLLM with fallback chain
    t0 = time.monotonic()
    all_models = [decision.model] + decision.fallbacks

    last_error = None
    for i, model_spec in enumerate(all_models):
        try:
            litellm_kwargs["model"] = model_spec.litellm_model
            _set_provider_keys(litellm_kwargs, model_spec)

            if stream:
                response = await litellm.acompletion(**litellm_kwargs)
                return StreamingResponse(
                    _stream_wrapper(response, model_spec, router, t0, decision),
                    media_type="text/event-stream",
                    headers={
                        "X-Model-Id": model_spec.id,
                        "X-Task-Tier": analysis.tier.value,
                        "X-Routing-Reason": decision.reason[:100],
                    },
                )
            else:
                response = await litellm.acompletion(**litellm_kwargs)
                elapsed = time.monotonic() - t0
                router.record_latency(model_spec.id, elapsed)

                # Extract response data
                resp_data = response.model_dump() if hasattr(response, "model_dump") else dict(response)

                # Calculate and record cost
                usage = resp_data.get("usage") or {}
                prompt_tokens = usage.get("prompt_tokens", 0) or 0
                completion_tokens = usage.get("completion_tokens", 0) or 0
                cost = _calculate_cost(model_spec, prompt_tokens, completion_tokens)
                router.record_cost(cost)

                # Add proxy metadata to response
                resp_data["_proxy"] = {
                    "model_id": model_spec.id,
                    "tier": analysis.tier.value,
                    "cost_usd": round(cost, 6),
                    "routing_reason": decision.reason,
                    "elapsed_ms": round(elapsed * 1000, 1),
                    "fallback_index": i,
                }

                return JSONResponse(
                    content=resp_data,
                    headers={
                        "X-Model-Id": model_spec.id,
                        "X-Task-Tier": analysis.tier.value,
                        "X-Cost-Usd": f"{cost:.6f}",
                        "X-Routing-Reason": decision.reason[:100],
                    },
                )

        except Exception as e:
            last_error = e
            logger.warning(
                "model_failed",
                model=model_spec.id,
                error=str(e)[:200],
                fallback_index=i,
            )
            continue

    # All models failed
    raise HTTPException(
        status_code=502,
        detail=f"All models failed. Last error: {str(last_error)[:200]}",
    )


async def _stream_wrapper(response, model_spec, router, t0, decision):
    """Wrap streaming response to track cost after completion."""
    prompt_tokens = 0
    completion_tokens = 0

    async for chunk in response:
        # Track tokens from stream
        if hasattr(chunk, "usage") and chunk.usage:
            prompt_tokens = getattr(chunk.usage, "prompt_tokens", 0) or 0
            completion_tokens = getattr(chunk.usage, "completion_tokens", 0) or 0

        yield f"data: {json.dumps(chunk.model_dump() if hasattr(chunk, 'model_dump') else dict(chunk))}\n\n"

    # Record cost after stream completes
    elapsed = time.monotonic() - t0
    router.record_latency(model_spec.id, elapsed)
    cost = _calculate_cost(model_spec, prompt_tokens, completion_tokens)
    router.record_cost(cost)

    yield "data: [DONE]\n\n"


def _calculate_cost(model: 'ModelSpec', prompt_tokens: int, completion_tokens: int) -> float:
    """Calculate actual cost from token counts."""
    input_cost = (prompt_tokens / 1_000_000) * model.input_cost
    output_cost = (completion_tokens / 1_000_000) * model.output_cost
    return input_cost + output_cost


def _set_provider_keys(kwargs: dict, model: 'ModelSpec') -> None:
    """Inject the correct API key for a model's provider."""
    settings = get_settings()
    provider_keys = {
        "anthropic": ("api_key", settings.anthropic_api_key),
        "openai": ("api_key", settings.openai_api_key),
        "openrouter": ("api_key", settings.openrouter_api_key),
        "google": ("api_key", settings.google_ai_key),
        "deepseek": ("api_key", settings.deepseek_api_key),
        "groq": ("api_key", settings.groq_api_key),
        "together": ("api_key", settings.together_api_key),
        "mistral": ("api_key", settings.mistral_api_key),
        "fireworks": ("api_key", settings.fireworks_api_key),
        "cerebras": ("api_key", settings.cerebras_api_key),
        "ollama": ("api_base", settings.ollama_base_url),
    }
    entry = provider_keys.get(model.provider)
    if entry:
        key_name, key_value = entry
        if key_value:
            kwargs[key_name] = key_value


# ── CLI entry point ──────────────────────────────────────────

def cli():
    parser = argparse.ArgumentParser(description="AI Proxy Server")
    parser.add_argument("--host", default=None, help="Bind host")
    parser.add_argument("--port", type=int, default=None, help="Bind port")
    parser.add_argument("--reload", action="store_true", help="Auto-reload on changes")
    args = parser.parse_args()

    s = get_settings()
    uvicorn.run(
        "proxym.main:app",
        host=args.host or s.host,
        port=args.port or s.port,
        reload=args.reload,
        log_level=s.log_level,
    )


if __name__ == "__main__":
    cli()
