"""
Security Middleware and Audit Logging

Provides:
- Request authentication and validation
- IP-based security controls
- Audit logging for compliance
- Request/response timing
- Security headers
"""
from __future__ import annotations

import hashlib
import time
import uuid
from typing import Any, Callable, Optional

from fastapi import HTTPException, Request, Response
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware

from app.auth.api_keys import ApiKeyInfo, validate_api_key, verify_request_signature
from app.auth.rate_limiter import (
    RateLimitError,
    ConcurrencyLimitError,
    check_rate_limit,
    get_remaining_quota,
)
from app.observability.logging import get_logger
from app.storage.redis import get_redis


_logger = get_logger("security")

# Paths that don't require API key authentication
PUBLIC_PATHS = {
    "/",
    "/health",
    "/ready",
    "/metrics",
    "/static",
    "/docs",
    "/openapi.json",
    "/redoc",
    "/api/v1/dashboard/auth/register",
    "/api/v1/dashboard/auth/login",
}

# Paths that use JWT auth instead of API key
JWT_AUTH_PATHS = {
    "/api/voice/session",
    "/voice",
}


def get_client_ip(request: Request) -> str:
    """Extract real client IP from request, handling proxies."""
    # Check common proxy headers
    forwarded = request.headers.get("x-forwarded-for")
    if forwarded:
        # Take first IP (original client)
        return forwarded.split(",")[0].strip()
    
    real_ip = request.headers.get("x-real-ip")
    if real_ip:
        return real_ip.strip()
    
    # Fallback to direct connection
    if request.client:
        return request.client.host
    
    return "unknown"


def get_request_id(request: Request) -> str:
    """Get or generate request ID for tracing."""
    request_id = request.headers.get("x-request-id")
    if not request_id:
        request_id = str(uuid.uuid4())
    return request_id


class SecurityMiddleware(BaseHTTPMiddleware):
    """
    Security middleware that handles:
    - API key validation
    - Rate limiting
    - Security headers
    - Request/response logging
    """
    
    def __init__(self, app, redis_factory: Callable):
        super().__init__(app)
        self.redis_factory = redis_factory
    
    async def dispatch(self, request: Request, call_next) -> Response:
        start_time = time.perf_counter()
        request_id = get_request_id(request)
        client_ip = get_client_ip(request)
        
        # Attach to request state for downstream use
        request.state.request_id = request_id
        request.state.client_ip = client_ip
        request.state.start_time = start_time
        
        path = request.url.path
        
        # Skip auth for public paths
        if self._is_public_path(path):
            response = await call_next(request)
            return self._add_security_headers(response, request_id, start_time)
        
        # Skip API key auth for JWT-authenticated paths
        if self._is_jwt_path(path):
            response = await call_next(request)
            return self._add_security_headers(response, request_id, start_time)
        
        # Validate API key
        api_key = self._extract_api_key(request)
        if not api_key:
            _logger.warning(
                "missing_api_key",
                path=path,
                client_ip=client_ip,
                request_id=request_id,
            )
            return self._error_response(401, "Missing API key", request_id)
        
        try:
            redis = await self.redis_factory()
            key_info = await validate_api_key(redis, api_key, client_ip)
            
            if not key_info:
                _logger.warning(
                    "invalid_api_key",
                    path=path,
                    client_ip=client_ip,
                    request_id=request_id,
                )
                return self._error_response(401, "Invalid API key", request_id)
            
            # Verify request signature if provided
            signature = request.headers.get("x-signature")
            timestamp_str = request.headers.get("x-timestamp")
            if signature and timestamp_str:
                try:
                    timestamp = int(timestamp_str)
                    body = await request.body()
                    if not verify_request_signature(api_key, timestamp, body, signature):
                        _logger.warning(
                            "invalid_signature",
                            path=path,
                            client_ip=client_ip,
                            request_id=request_id,
                        )
                        return self._error_response(401, "Invalid request signature", request_id)
                except (ValueError, TypeError):
                    return self._error_response(400, "Invalid timestamp", request_id)
            
            # Check rate limits
            try:
                await check_rate_limit(redis, key_info.org_id, key_info.limits, path)
            except RateLimitError as e:
                _logger.warning(
                    "rate_limit_exceeded",
                    org_id=key_info.org_id,
                    limit_type=e.limit_type,
                    limit=e.limit,
                    current=e.current,
                    request_id=request_id,
                )
                return self._rate_limit_response(e, request_id)
            
            # Attach key info to request state
            request.state.api_key_info = key_info
            request.state.org_id = key_info.org_id
            
            # Get remaining quota for headers
            quota_headers = await get_remaining_quota(redis, key_info.org_id, key_info.limits)
            
            # Process request
            response = await call_next(request)
            
            # Add rate limit headers
            for header, value in quota_headers.items():
                response.headers[header] = str(value)
            
            # Audit log successful request
            self._audit_log(
                request_id=request_id,
                org_id=key_info.org_id,
                path=path,
                method=request.method,
                client_ip=client_ip,
                status_code=response.status_code,
                duration_ms=int((time.perf_counter() - start_time) * 1000),
            )
            
            return self._add_security_headers(response, request_id, start_time)
            
        except Exception as e:
            _logger.error(
                "security_middleware_error",
                error=str(e),
                path=path,
                request_id=request_id,
            )
            return self._error_response(500, "Internal security error", request_id)
    
    def _is_public_path(self, path: str) -> bool:
        """Check if path is public (no auth required)."""
        if path in PUBLIC_PATHS:
            return True
        for public in PUBLIC_PATHS:
            if path.startswith(public + "/"):
                return True
        return False
    
    def _is_jwt_path(self, path: str) -> bool:
        """Check if path uses JWT auth instead of API key."""
        if path in JWT_AUTH_PATHS:
            return True
        # All dashboard endpoints except auth use JWT
        if path.startswith("/api/v1/dashboard") and not path.startswith("/api/v1/dashboard/auth"):
            return True
        if path.startswith("/voice"):
            return True
        return False
    
    def _extract_api_key(self, request: Request) -> Optional[str]:
        """Extract API key from request."""
        # Check Authorization header
        auth = request.headers.get("authorization", "")
        if auth.startswith("Bearer "):
            token = auth[7:].strip()
            if token.startswith("sv_"):
                return token
        
        # Check X-API-Key header
        api_key = request.headers.get("x-api-key")
        if api_key:
            return api_key
        
        # Check query parameter (not recommended but supported)
        api_key = request.query_params.get("api_key")
        if api_key:
            return api_key
        
        return None
    
    def _error_response(self, status: int, message: str, request_id: str) -> JSONResponse:
        """Create error response with security headers."""
        response = JSONResponse(
            status_code=status,
            content={"error": message, "request_id": request_id},
        )
        return self._add_security_headers(response, request_id, time.perf_counter())
    
    def _rate_limit_response(self, error: RateLimitError, request_id: str) -> JSONResponse:
        """Create rate limit error response."""
        response = JSONResponse(
            status_code=429,
            content={
                "error": str(error),
                "limit_type": error.limit_type,
                "limit": error.limit,
                "current": error.current,
                "retry_after": error.retry_after,
                "request_id": request_id,
            },
        )
        response.headers["Retry-After"] = str(error.retry_after)
        return self._add_security_headers(response, request_id, time.perf_counter())
    
    def _add_security_headers(
        self,
        response: Response,
        request_id: str,
        start_time: float,
    ) -> Response:
        """Add security headers to response."""
        response.headers["X-Request-ID"] = request_id
        response.headers["X-Response-Time"] = f"{int((time.perf_counter() - start_time) * 1000)}ms"
        response.headers["X-Content-Type-Options"] = "nosniff"
        response.headers["X-Frame-Options"] = "DENY"
        response.headers["X-XSS-Protection"] = "1; mode=block"
        response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"
        response.headers["Cache-Control"] = "no-store, no-cache, must-revalidate"
        response.headers["Pragma"] = "no-cache"
        return response
    
    def _audit_log(
        self,
        request_id: str,
        org_id: str,
        path: str,
        method: str,
        client_ip: str,
        status_code: int,
        duration_ms: int,
    ) -> None:
        """Write audit log entry."""
        _logger.info(
            "api_request",
            request_id=request_id,
            org_id=org_id,
            path=path,
            method=method,
            client_ip=client_ip,
            status_code=status_code,
            duration_ms=duration_ms,
        )


class IPAllowlistMiddleware(BaseHTTPMiddleware):
    """
    Optional middleware for IP allowlisting.
    
    Enable by setting ALLOWED_IPS environment variable.
    """
    
    def __init__(self, app, allowed_ips: Optional[list[str]] = None):
        super().__init__(app)
        self.allowed_ips = set(allowed_ips or [])
        self.enabled = bool(self.allowed_ips)
    
    async def dispatch(self, request: Request, call_next) -> Response:
        if not self.enabled:
            return await call_next(request)
        
        client_ip = get_client_ip(request)
        
        # Allow localhost in development
        if client_ip in ("127.0.0.1", "::1", "localhost"):
            return await call_next(request)
        
        if client_ip not in self.allowed_ips:
            _logger.warning(
                "ip_blocked",
                client_ip=client_ip,
                path=request.url.path,
            )
            return JSONResponse(
                status_code=403,
                content={"error": "Access denied"},
            )
        
        return await call_next(request)


def mask_sensitive_data(data: dict, fields: list[str] = None) -> dict:
    """Mask sensitive fields in data for logging."""
    if fields is None:
        fields = ["api_key", "password", "token", "secret", "authorization"]
    
    masked = {}
    for key, value in data.items():
        if any(f in key.lower() for f in fields):
            if isinstance(value, str) and len(value) > 8:
                masked[key] = value[:4] + "****" + value[-4:]
            else:
                masked[key] = "****"
        elif isinstance(value, dict):
            masked[key] = mask_sensitive_data(value, fields)
        else:
            masked[key] = value
    return masked
