"""Server middleware - Authentication and Rate Limiting"""

import hmac
import os
import time
from collections import defaultdict
from pathlib import Path
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from starlette.responses import JSONResponse


def _load_client_keys(config_dir: str | None = None) -> dict[str, str]:
    """Load client API keys from clients.yaml.

    Returns a dict mapping client_name -> api_key.
    """
    if config_dir is None:
        # Try current directory and home directory
        for candidate in [Path.cwd(), Path.home() / ".agentkit"]:
            clients_path = candidate / "clients.yaml"
            if clients_path.exists():
                config_dir = str(candidate)
                break
        else:
            return {}

    import yaml
    clients_path = Path(config_dir) / "clients.yaml"
    if not clients_path.exists():
        return {}

    with open(clients_path, encoding="utf-8") as f:
        data = yaml.safe_load(f) or {}

    # data is {client_name: {api_key: "...", ...}}
    return {name: info["api_key"] for name, info in data.items() if "api_key" in info}


class APIKeyAuthMiddleware(BaseHTTPMiddleware):
    """API Key authentication middleware.

    Validates X-API-Key header against:
    1. api_key parameter (global key, passed directly)
    2. Client keys from clients.yaml (generated by `agentkit pair`)

    Skips validation if no keys are configured (dev mode).
    Whitelisted paths (no auth required): /api/v1/health
    """

    WHITELIST_PATHS = ("/api/v1/health",)

    def __init__(self, app, api_key: str | None = None):
        super().__init__(app)
        self._api_key = api_key

    async def dispatch(self, request: Request, call_next):
        # Skip auth for whitelisted paths
        if any(request.url.path.startswith(p) for p in self.WHITELIST_PATHS):
            return await call_next(request)

        # Collect all valid keys
        valid_keys = set()

        # Global key from parameter
        if self._api_key:
            valid_keys.add(self._api_key)

        # Client keys from clients.yaml
        client_keys = _load_client_keys()
        valid_keys.update(client_keys.values())

        # No keys configured = dev mode
        if not valid_keys:
            return await call_next(request)

        # Check API key from header
        provided_key = request.headers.get("X-API-Key")
        if not provided_key or not any(hmac.compare_digest(provided_key.encode(), k.encode()) for k in valid_keys):
            return JSONResponse(
                status_code=401,
                content={"error": "Unauthorized", "message": "Invalid or missing API key"},
            )

        return await call_next(request)


class RateLimiter:
    """Fixed-window rate limiter.
    
    Tracks request counts per key (IP or API key) within time windows.
    """
    
    def __init__(self, max_requests: int = 60, window_seconds: int = 60):
        self._max_requests = max_requests
        self._window_seconds = window_seconds
        self._requests: dict[str, list[float]] = defaultdict(list)
    
    def is_allowed(self, key: str) -> tuple[bool, float]:
        """Check if request is allowed. Returns (allowed, retry_after_seconds)."""
        now = time.time()
        window_start = now - self._window_seconds
        
        # Clean old requests outside the window
        self._requests[key] = [
            ts for ts in self._requests[key] if ts > window_start
        ]
        
        if len(self._requests[key]) >= self._max_requests:
            retry_after = self._requests[key][0] + self._window_seconds - now
            return False, max(0, retry_after)
        
        self._requests[key].append(now)
        return True, 0.0
    
    @property
    def max_requests(self) -> int:
        return self._max_requests


class RateLimitMiddleware(BaseHTTPMiddleware):
    """Rate limiting middleware.
    
    Limits requests per IP. Returns 429 Too Many Requests when exceeded.
    Configurable via AGENTKIT_RATE_LIMIT_PER_MINUTE env var (default: 60).
    """
    
    def __init__(self, app, max_requests: int | None = None, window_seconds: int = 60):
        super().__init__(app)
        if max_requests is None:
            max_requests = int(os.environ.get("AGENTKIT_RATE_LIMIT_PER_MINUTE", "60"))
        self._limiter = RateLimiter(max_requests=max_requests, window_seconds=window_seconds)
    
    async def dispatch(self, request: Request, call_next):
        # Use API key if available, otherwise IP
        api_key = request.headers.get("X-API-Key")
        key = f"key:{api_key}" if api_key else f"ip:{request.client.host}"
        
        allowed, retry_after = self._limiter.is_allowed(key)
        if not allowed:
            return JSONResponse(
                status_code=429,
                content={
                    "error": "Too Many Requests",
                    "message": f"Rate limit exceeded. Try again in {int(retry_after)} seconds.",
                },
                headers={"Retry-After": str(int(retry_after))},
            )
        
        response = await call_next(request)
        return response
