"""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, Response
from fastapi.staticfiles import StaticFiles
from fastapi.middleware.cors import CORSMiddleware

from proxym.cache import DeltaBuffer
from proxym.config import get_settings
from proxym.context import detect_project, ProjectContext, SessionStore
from proxym.context.tool_selector import select_tools
from proxym.mcp import MCPRegistry, MCPRouter
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
_mcp_registry: MCPRegistry | None = None
_mcp_router: MCPRouter | None = None
_session_store: SessionStore | None = None


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


def _get_mcp_router() -> MCPRouter:
    global _mcp_registry, _mcp_router
    if _mcp_router is None:
        _mcp_registry = MCPRegistry()
        _mcp_router = MCPRouter(_mcp_registry)
    return _mcp_router


def _get_session_store() -> SessionStore:
    global _session_store
    if _session_store is None:
        _session_store = SessionStore()
    return _session_store


# ── 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("/favicon.ico")
async def favicon():
    """Return empty favicon to prevent 401 errors."""
    return Response(content="", media_type="image/x-icon")


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


@app.get("/dashboard-ui")
async def dashboard_ui():
    """Serve the React dashboard UI."""
    # Try local development path first, then Docker container path
    paths_to_try = [
        "frontend/dashboard.html",  # Local development
        "/app/frontend/dashboard.html",  # Docker container
    ]
    
    for file_path in paths_to_try:
        try:
            with open(file_path, "r") as f:
                content = f.read()
            return Response(
                content=content,
                media_type="text/html",
                headers={"Content-Type": "text/html; charset=utf-8"}
            )
        except FileNotFoundError:
            continue
        except Exception as e:
            logger.error("dashboard_ui_error", error=str(e), path=file_path)
            continue
    
    # If neither path works, return an error message
    logger.error("dashboard_ui_not_found", paths=paths_to_try)
    return Response(
        content="<h1>Dashboard Error</h1><p>Could not load dashboard: HTML file not found</p>",
        status_code=500,
        media_type="text/html"
    )


@app.get("/proxym-dashboard.jsx")
async def dashboard_jsx():
    """Serve the React dashboard JSX file."""
    # Try local development path first, then Docker container path
    paths_to_try = [
        "frontend/proxym-dashboard.jsx",  # Local development
        "/app/frontend/proxym-dashboard.jsx",  # Docker container
    ]
    
    for file_path in paths_to_try:
        try:
            with open(file_path, "r") as f:
                content = f.read()
            return Response(
                content=content,
                media_type="text/javascript",
                headers={"Content-Type": "text/javascript; charset=utf-8"}
            )
        except FileNotFoundError:
            continue
        except Exception as e:
            logger.error("dashboard_jsx_error", error=str(e), path=file_path)
            continue
    
    # If neither path works, return a helpful error message
    logger.error("dashboard_jsx_not_found", paths=paths_to_try)
    return Response(
        content="// Dashboard JSX file not found\nconsole.error('Dashboard JSX file not found. Please ensure proxym-dashboard.jsx exists in the project root.');\n",
        media_type="text/javascript"
    )


@app.get("/config")
async def get_frontend_config():
    """Return configuration for frontend."""
    settings = get_settings()
    return {
        "api_url": settings.dashboard_api_url,
        "api_key": settings.default_api_key,
        "debug": settings.debug,
    }


@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"),
    }


# ── Context & session endpoints ───────────────────────────────

@app.post("/v1/context/detect")
async def detect_context(request: Request):
    """Detect project context from request body (messages + headers)."""
    body = await request.json()
    messages = body.get("messages", [])
    headers = {k.lower(): v for k, v in request.headers.items()}
    ctx = detect_project(messages, headers=headers)
    return {
        "project_id": ctx.project_id,
        "project_name": ctx.project_name,
        "project_root": ctx.project_root,
        "stack": ctx.stack,
        "git_remote": ctx.git_remote,
    }


@app.get("/v1/sessions")
async def list_sessions():
    """List active agent sessions."""
    store = _get_session_store()
    return {
        "sessions": [
            {
                "session_id": s.session_id,
                "project_id": s.project_id,
                "age_seconds": round(s.age_seconds, 1),
                "idle_seconds": round(s.idle_seconds, 1),
                "history_length": len(s.compressed_history),
                "active_files": s.active_files,
            }
            for s in store.list_sessions()
        ],
    }


@app.get("/v1/mcp/servers")
async def list_mcp_servers():
    """List registered MCP servers."""
    mcp_router = _get_mcp_router()
    servers = mcp_router.registry.list_servers(enabled_only=False)
    return {
        "servers": [
            {
                "name": s.name,
                "url": s.url,
                "transport": s.transport.value,
                "tools": s.tools,
                "enabled": s.enabled,
                "priority": s.priority,
                "description": s.description,
            }
            for s in servers
        ],
    }


# ── 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)
      X-Working-Directory: /path/to/project  (project detection)
      X-Session-Id: abc123  (session continuity)
      X-Enable-Tools: true  (inject MCP tools based on context)
    """
    body = await request.json()
    messages: list[dict] = body.get("messages", [])
    requested_model: str | None = body.get("model")
    stream: bool = body.get("stream", False)

    # 1. Inject delta context if requested
    if request.headers.get("X-Inject-Context", "").lower() == "true":
        _inject_context(messages)

    # 2. Detect project context
    headers = {k.lower(): v for k, v in request.headers.items()}
    project_ctx = detect_project(messages, headers=headers)

    # 3. Analyze + route
    analysis = analyze_request(
        messages,
        explicit_tier=x_task_tier,
        explicit_model=requested_model if requested_model and "/" in requested_model else None,
    )
    router = _get_router()
    explicit = _resolve_model_alias(requested_model)
    decision = router.route(analysis, explicit_model=explicit)

    # 4. Inject MCP tools if enabled and not already provided
    if headers.get("x-enable-tools", "").lower() == "true" and "tools" not in body:
        mcp_router = _get_mcp_router()
        servers = mcp_router.select_servers(project_ctx, analysis)
        if servers:
            mcp_tools = mcp_router.get_tools_for_servers(servers)
            if mcp_tools:
                body["tools"] = mcp_tools
                project_ctx.active_mcp_servers = [s.name for s in servers]

    # 5. Touch session
    session_id = headers.get("x-session-id", "")
    if session_id or project_ctx.project_id:
        store = _get_session_store()
        session = store.get_or_create(
            session_id=session_id, project_id=project_ctx.project_id,
        )
        project_ctx.session_id = session.session_id

    # 6. Build kwargs + execute with fallback chain
    litellm_kwargs = _build_litellm_kwargs(body, messages, stream, decision)
    return await _execute_with_fallbacks(
        litellm_kwargs, decision, analysis, router, stream,
    )


# ── chat_completions helpers ─────────────────────────────────

def _inject_context(messages: list[dict]) -> None:
    """Inject delta context payload into messages (mutates in place)."""
    if not _context_payload:
        return
    if messages and messages[0].get("role") == "system":
        messages[0]["content"] += f"\n\n{_context_payload}"
    else:
        messages.insert(0, {
            "role": "system",
            "content": f"Current codebase context:\n{_context_payload}",
        })


_TIER_ALIASES: dict[str, str] = {
    "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",
}


def _resolve_model_alias(requested_model: str | None) -> str | None:
    """Resolve tier aliases (cheap, balanced, etc.) to model IDs."""
    if not requested_model:
        return None
    return _TIER_ALIASES.get(requested_model, requested_model)


_PASSTHROUGH_PARAMS = (
    "temperature", "max_tokens", "top_p", "stop", "tools",
    "tool_choice", "response_format", "seed",
)


def _build_litellm_kwargs(
    body: dict, messages: list[dict], stream: bool, decision: RoutingDecision,
) -> dict:
    """Build LiteLLM completion kwargs from request body and routing decision."""
    kwargs = {
        "model": decision.model.litellm_model,
        "messages": messages,
        "stream": stream,
    }
    for key in _PASSTHROUGH_PARAMS:
        if key in body:
            kwargs[key] = body[key]
    _set_provider_keys(kwargs, decision.model)
    return kwargs


async def _execute_with_fallbacks(
    litellm_kwargs: dict,
    decision: RoutingDecision,
    analysis: AnalysisResult,
    router: Router,
    stream: bool,
) -> StreamingResponse | JSONResponse:
    """Try primary model + fallbacks, return first successful response."""
    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:
                return await _handle_stream(
                    litellm_kwargs, model_spec, analysis, router, t0, decision,
                )
            else:
                return await _handle_non_stream(
                    litellm_kwargs, model_spec, analysis, router, t0, decision, i,
                )
        except Exception as e:
            last_error = e
            logger.warning(
                "model_failed",
                model=model_spec.id,
                error=str(e)[:200],
                fallback_index=i,
            )
            continue

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


async def _handle_stream(
    litellm_kwargs, model_spec, analysis, router, t0, decision,
) -> StreamingResponse:
    """Execute streaming completion and return StreamingResponse."""
    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],
        },
    )


async def _handle_non_stream(
    litellm_kwargs, model_spec, analysis, router, t0, decision, fallback_index,
) -> JSONResponse:
    """Execute non-streaming completion, record cost, return JSONResponse."""
    response = await litellm.acompletion(**litellm_kwargs)
    elapsed = time.monotonic() - t0
    router.record_latency(model_spec.id, elapsed)

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

    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)

    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": fallback_index,
    }

    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],
        },
    )


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()
