"""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

Endpoints are split into sub-routers under proxym.api/:
  - completions.py  — /v1/chat/completions + helpers
  - system.py       — /health, /v1/models, /v1/budget, /favicon.ico
  - context_api.py  — /v1/context/*, /v1/sessions, /v1/mcp/servers
  - frontend.py     — /dashboard-ui, /proxym-dashboard.jsx, /config, JS assets
"""

from __future__ import annotations

import argparse
import logging
import logging.handlers
from contextlib import asynccontextmanager
from pathlib import Path

import litellm
import structlog
import uvicorn
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from proxym.api import (
    completions_router,
    console_router,
    context_router,
    frontend_router,
    system_router,
    tests_router,
    vscode_router,
    voice_actions_router,
)
from proxym.config import get_settings
from proxym.middleware import AuthMiddleware, CostTrackingMiddleware
from proxym.providers import MODELS

# ── Logging: console + file ──────────────────────────────────
_LOG_DIR = Path.home() / ".local" / "share" / "proxym"
_LOG_DIR.mkdir(parents=True, exist_ok=True)
_LOG_FILE = _LOG_DIR / "proxym.log"

_file_handler = logging.handlers.RotatingFileHandler(
    _LOG_FILE, maxBytes=5_000_000, backupCount=3, encoding="utf-8",
)
_file_handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s"))

_console_handler = logging.StreamHandler()
_console_handler.setLevel(logging.DEBUG)

logging.basicConfig(level=logging.INFO, handlers=[_console_handler, _file_handler])

structlog.configure(
    processors=[
        structlog.stdlib.add_log_level,
        structlog.stdlib.ProcessorFormatter.wrap_for_formatter,
    ],
    wrapper_class=structlog.stdlib.BoundLogger,
    context_class=dict,
    logger_factory=structlog.stdlib.LoggerFactory(),
)

# Attach structlog renderer to both handlers
_formatter = structlog.stdlib.ProcessorFormatter(
    processor=structlog.dev.ConsoleRenderer(),
)
_console_handler.setFormatter(_formatter)
_file_handler.setFormatter(_formatter)

logger = structlog.get_logger()


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

    # Export provider API keys to os.environ so litellm can discover them.
    # Pydantic-settings loads .env into the Settings object but does NOT
    # export to os.environ — litellm reads env vars directly.
    import os
    _env_keys = {
        "ANTHROPIC_API_KEY": settings.anthropic_api_key,
        "OPENAI_API_KEY": settings.openai_api_key,
        "OPENROUTER_API_KEY": settings.openrouter_api_key,
        "GOOGLE_AI_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,
    }
    for env_name, env_val in _env_keys.items():
        if env_val and env_name not in os.environ:
            os.environ[env_name] = env_val

    # Also set litellm-specific key attributes
    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

    # Start tool executor (processes queued jobs)
    from proxym.tools.executor import ToolExecutor
    tool_executor = ToolExecutor.get()
    await tool_executor.start()
    app.state.tool_executor = tool_executor  # type: ignore[attr-defined]

    # Start health watchdog
    if settings.watchdog_enabled:
        from proxym.observability.watchdog import HealthWatchdog

        repair_callback = None
        if settings.repair_enabled and settings.repair_mode == "auto":
            from proxym.observability.repair_agent import RepairAgent
            agent = RepairAgent(mode="auto")
            repair_callback = agent.diagnose_and_repair

        watchdog = HealthWatchdog(
            interval_seconds=settings.watchdog_interval,
            repair_callback=repair_callback,
        )
        await watchdog.start()
        app.state.watchdog = watchdog  # type: ignore[attr-defined]

    yield

    # Shutdown tool executor
    if hasattr(app.state, "tool_executor"):
        await app.state.tool_executor.stop()

    # Shutdown watchdog
    if hasattr(app.state, "watchdog"):
        await app.state.watchdog.stop()

    logger.info("proxy_stopped")


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

# Mount sub-routers
from proxym.dashboard import router as dashboard_router
from proxym.mcp.self_server import router as self_mcp_router
from proxym.api.observability_api import router as obs_router
app.include_router(dashboard_router)
app.include_router(self_mcp_router)
app.include_router(obs_router)
app.include_router(system_router)
app.include_router(frontend_router)
app.include_router(context_router)
app.include_router(completions_router)
app.include_router(console_router)
app.include_router(tests_router)
app.include_router(vscode_router)
app.include_router(voice_actions_router)

# Add middleware
settings = get_settings()

# CORS middleware - allow frontend on different port
app.add_middleware(
    CORSMiddleware,
    allow_origins=["http://localhost:4000", "http://localhost:4001", "http://127.0.0.1:4000", "http://127.0.0.1:4001"],
    allow_credentials=True,
    allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
    allow_headers=["*"],
)

app.add_middleware(CostTrackingMiddleware)
app.add_middleware(AuthMiddleware, master_key=settings.master_key)


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