from fastapi import Request, HTTPException, status
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse
from typing import List
import ipaddress
from ...shared.config.settings import settings
from ...shared.config.logging_config import get_logger

logger = get_logger(__name__)

class AdminIPWhitelistMiddleware(BaseHTTPMiddleware):
    """
    Middleware to enforce IP whitelisting for admin endpoints.
    Only applies to routes starting with /api/admin/ or containing 'admin' in path.
    """
    
    def __init__(self, app, enabled: bool = None, whitelist: List[str] = None):
        super().__init__(app)
        self.enabled = enabled if enabled is not None else settings.IP_WHITELIST_ENABLED
        self.whitelist = whitelist if whitelist is not None else settings.ADMIN_IP_WHITELIST
        
        # Pre-compile IP networks for faster lookup
        self._compiled_networks = []
        if self.enabled and self.whitelist:
            for ip_or_cidr in self.whitelist:
                try:
                    if '/' in ip_or_cidr:
                        # CIDR notation
                        network = ipaddress.ip_network(ip_or_cidr, strict=False)
                        self._compiled_networks.append(network)
                    else:
                        # Single IP address
                        ip = ipaddress.ip_address(ip_or_cidr)
                        # Convert to /32 network for consistent handling
                        self._compiled_networks.append(ipaddress.ip_network(f'{ip}/32', strict=False))
                except (ValueError, ipaddress.AddressValueError) as e:
                    logger.warning(f'Invalid IP/CIDR in admin whitelist: {ip_or_cidr} - {str(e)}')
        
        if self.enabled:
            logger.info(f'Admin IP whitelisting enabled with {len(self._compiled_networks)} allowed IP(s)/CIDR range(s)')
    
    def _is_admin_route(self, path: str) -> bool:
        """Check if the path is an admin route"""
        return '/admin/' in path.lower() or path.lower().startswith('/api/admin')
    
    def _get_client_ip(self, request: Request) -> str:
        """Extract client IP address from request"""
        # Check for forwarded IP (when behind proxy/load balancer)
        forwarded_for = request.headers.get("X-Forwarded-For")
        if forwarded_for:
            # X-Forwarded-For can contain multiple IPs, take the first one (original client)
            return forwarded_for.split(",")[0].strip()
        
        real_ip = request.headers.get("X-Real-IP")
        if real_ip:
            return real_ip.strip()
        
        # Fallback to direct client IP
        if request.client:
            return request.client.host
        
        return None
    
    def _is_ip_allowed(self, ip_address: str) -> bool:
        """Check if IP address is in whitelist"""
        if not self._compiled_networks:
            # Empty whitelist means deny all (security-first approach)
            return False
        
        try:
            client_ip = ipaddress.ip_address(ip_address)
            for network in self._compiled_networks:
                if client_ip in network:
                    return True
            return False
        except (ValueError, ipaddress.AddressValueError):
            logger.warning(f'Invalid IP address format: {ip_address}')
            return False
    
    async def dispatch(self, request: Request, call_next):
        # Skip if not enabled
        if not self.enabled:
            return await call_next(request)
        
        # Skip OPTIONS requests (CORS preflight) - let CORS middleware handle them
        if request.method == 'OPTIONS':
            return await call_next(request)
        
        # Only apply to admin routes
        if not self._is_admin_route(request.url.path):
            return await call_next(request)
        
        # Skip IP check for health checks and public endpoints
        if request.url.path in ['/health', '/api/health', '/metrics']:
            return await call_next(request)
        
        client_ip = self._get_client_ip(request)
        
        if not client_ip:
            logger.warning("Could not determine client IP address for admin route")
            # Deny by default if IP cannot be determined (security-first)
            return JSONResponse(
                status_code=status.HTTP_403_FORBIDDEN,
                content={"status": "error", "message": "Access denied: Unable to verify IP address"}
            )
        
        # Check whitelist
        if not self._is_ip_allowed(client_ip):
            logger.warning(
                f"Admin route access denied for IP: {client_ip} from path: {request.url.path}"
            )
            return JSONResponse(
                status_code=status.HTTP_403_FORBIDDEN,
                content={"status": "error", "message": "Access denied. IP address not whitelisted."}
            )
        
        # IP is whitelisted, continue
        return await call_next(request)
