from typing import Dict, List, Optional, Set
from pydantic import BaseModel
from enum import Enum

class IPControlMode(str, Enum):
    ALLOWLIST = "allowlist"
    BLOCKLIST = "blocklist"
    DISABLED = "disabled"

class IPRule(BaseModel):
    """Individual IP access rule."""
    name: str
    description: Optional[str] = None
    mode: IPControlMode
    ips: List[str] = []  # IPs and CIDR ranges
    countries: List[str] = []  # ISO country codes
    applies_to_paths: List[str] = ["*"]  # Path patterns
    applies_to_roles: List[str] = ["*"]  # User roles
    applies_to_users: List[str] = []  # Specific user IDs
    priority: int = 0  # Higher priority rules evaluated first
    enabled: bool = True
    expires_at: Optional[datetime] = None

class IPControlConfig(BaseModel):
    """IP access control configuration."""
    # Global settings
    enabled: bool = True
    default_mode: IPControlMode = IPControlMode.DISABLED
    check_behind_proxy: bool = True
    trusted_proxies: List[str] = []
    enforcement_level: str = "enforce"  # enforce, monitor, off
    
    # Rules
    rules: List[IPRule] = []
    
    # Performance settings
    cache_ttl: int = 300  # Cache IP checks for 5 minutes
    use_geoip: bool = False
    geoip_database: Optional[str] = None
    
    # Emergency access
    bypass_enabled: bool = False
    bypass_code_ttl: int = 300
    
    # Integration
    use_ip_reputation: bool = False
    ip_reputation_threshold: float = 0.5

# Add to main config
ip_control: IPControlConfig = IPControlConfig()

def upgrade():
    op.create_table(
        'ip_rules',
        sa.Column('id', sa.String, primary_key=True),
        sa.Column('name', sa.String, nullable=False, unique=True),
        sa.Column('description', sa.Text),
        sa.Column('mode', sa.String, nullable=False),  # allowlist/blocklist
        sa.Column('rule_data', sa.JSON, nullable=False),  # Full rule config
        sa.Column('enabled', sa.Boolean, default=True),
        sa.Column('priority', sa.Integer, default=0),
        sa.Column('expires_at', sa.DateTime(timezone=True)),
        sa.Column('created_at', sa.DateTime(timezone=True), default=func.now()),
        sa.Column('created_by', sa.String, sa.ForeignKey('users.id')),
        sa.Column('updated_at', sa.DateTime(timezone=True), onupdate=func.now()),
        sa.Column('updated_by', sa.String, sa.ForeignKey('users.id'))
    )
    
    op.create_index('idx_ip_rules_enabled_priority', 'ip_rules', 
                    ['enabled', 'priority'])
    
    # Temporary IP blocks table
    op.create_table(
        'ip_blocks',
        sa.Column('id', sa.String, primary_key=True),
        sa.Column('ip_address', sa.String, nullable=False),
        sa.Column('reason', sa.Text),
        sa.Column('blocked_until', sa.DateTime(timezone=True), nullable=False),
        sa.Column('blocked_by', sa.String, sa.ForeignKey('users.id')),
        sa.Column('created_at', sa.DateTime(timezone=True), default=func.now())
    )
    
    op.create_index('idx_ip_blocks_ip_until', 'ip_blocks', 
                    ['ip_address', 'blocked_until'])

import re
import ipaddress
from datetime import datetime, timezone
from typing import Optional, Tuple, List
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.requests import Request
from fastapi import HTTPException

from mcpgateway.config import settings
from mcpgateway.services.ip_control_service import IPControlService
from mcpgateway.services.audit_logger import audit_logger
from mcpgateway.auth import get_current_user_optional

class IPControlMiddleware(BaseHTTPMiddleware):
    """Generic IP-based access control middleware."""
    
    def __init__(self, app):
        super().__init__(app)
        self.ip_service = IPControlService()
        self._path_cache = {}  # Cache compiled regex patterns
    
    async def dispatch(self, request: Request, call_next):
        if not settings.ip_control.enabled:
            return await call_next(request)
        
        # Extract client IP
        client_ip = self._get_client_ip(request)
        request.state.client_ip = client_ip
        
        # Get current user context
        user = await get_current_user_optional(request)
        user_id = user.id if user else None
        user_roles = user.roles if user else []
        
        # Check IP access
        allowed, rule_name, reason = await self.ip_service.check_access(
            ip=client_ip,
            path=request.url.path,
            user_id=user_id,
            user_roles=user_roles
        )
        
        # Log the access attempt
        await self._log_access(
            request=request,
            client_ip=client_ip,
            allowed=allowed,
            rule_name=rule_name,
            reason=reason,
            user_id=user_id
        )
        
        # Store result in request state
        request.state.ip_check = {
            "allowed": allowed,
            "rule": rule_name,
            "reason": reason
        }
        
        # Enforce based on configuration
        if not allowed and settings.ip_control.enforcement_level == "enforce":
            raise HTTPException(
                status_code=403,
                detail="Access denied from your IP address"
            )
        
        response = await call_next(request)
        
        # Add security headers
        if not allowed and settings.ip_control.enforcement_level == "monitor":
            response.headers["X-IP-Access-Warning"] = "monitored"
        
        return response
    
    def _get_client_ip(self, request: Request) -> str:
        """Extract real client IP considering proxies."""
        if settings.ip_control.check_behind_proxy:
            forwarded_for = request.headers.get("X-Forwarded-For")
            real_ip = request.headers.get("X-Real-IP")
            
            if forwarded_for and self._is_trusted_proxy(request.client.host):
                # Take the first IP from X-Forwarded-For chain
                return forwarded_for.split(",")[0].strip()
            elif real_ip and self._is_trusted_proxy(request.client.host):
                return real_ip.strip()
        
        return request.client.host
    
    def _is_trusted_proxy(self, proxy_ip: str) -> bool:
        """Check if proxy IP is trusted."""
        return any(
            self._ip_matches(proxy_ip, trusted)
            for trusted in settings.ip_control.trusted_proxies
        )

class IPControlService:
    """Service for IP access control logic."""
    
    def __init__(self):
        self.redis = get_redis_client()
        self.geoip = GeoIPService() if settings.ip_control.use_geoip else None
        self.reputation = IPReputationService() if settings.ip_control.use_ip_reputation else None
    
    async def check_access(
        self,
        ip: str,
        path: str,
        user_id: Optional[str] = None,
        user_roles: List[str] = []
    ) -> Tuple[bool, Optional[str], str]:
        """
        Check if IP has access to the given path.
        Returns: (allowed, rule_name, reason)
        """
        # Check cache first
        cache_key = f"ip_access:{ip}:{path}:{user_id}:{':'.join(sorted(user_roles))}"
        cached = await self.redis.get(cache_key)
        if cached:
            return json.loads(cached)
        
        # Check temporary blocks
        if await self._is_temporarily_blocked(ip):
            result = (False, "temporary_block", "IP is temporarily blocked")
            await self._cache_result(cache_key, result)
            return result
        
        # Check IP reputation if enabled
        if self.reputation:
            reputation_score = await self.reputation.get_score(ip)
            if reputation_score < settings.ip_control.ip_reputation_threshold:
                result = (False, "reputation", f"IP reputation too low: {reputation_score}")
                await self._cache_result(cache_key, result)
                return result
        
        # Get applicable rules
        rules = await self._get_applicable_rules(path, user_id, user_roles)
        
        # Evaluate rules in priority order
        for rule in sorted(rules, key=lambda r: r.priority, reverse=True):
            matches = await self._evaluate_rule(rule, ip, path, user_id, user_roles)
            
            if matches:
                if rule.mode == IPControlMode.ALLOWLIST:
                    result = (True, rule.name, "Allowed by allowlist rule")
                else:  # BLOCKLIST
                    result = (False, rule.name, "Blocked by blocklist rule")
                
                await self._cache_result(cache_key, result)
                return result
        
        # No matching rules - use default mode
        if settings.ip_control.default_mode == IPControlMode.ALLOWLIST:
            result = (False, "default", "Not in any allowlist")
        elif settings.ip_control.default_mode == IPControlMode.BLOCKLIST:
            result = (True, "default", "Not in any blocklist")
        else:  # DISABLED
            result = (True, "default", "IP control disabled")
        
        await self._cache_result(cache_key, result)
        return result
    
    async def _evaluate_rule(
        self,
        rule: IPRule,
        ip: str,
        path: str,
        user_id: Optional[str],
        user_roles: List[str]
    ) -> bool:
        """Check if a rule matches the current request."""
        # Check if rule applies to this path
        if not self._path_matches(path, rule.applies_to_paths):
            return False
        
        # Check if rule applies to this user/role
        if rule.applies_to_users and user_id not in rule.applies_to_users:
            return False
        
        if rule.applies_to_roles != ["*"]:
            if not any(role in rule.applies_to_roles for role in user_roles):
                return False
        
        # Check IP matches
        if rule.ips:
            if any(self._ip_matches(ip, allowed_ip) for allowed_ip in rule.ips):
                return True
        
        # Check country matches
        if rule.countries and self.geoip:
            country = await self.geoip.get_country(ip)
            if country in rule.countries:
                return True
        
        return False

@router.get("/api/ip-control/rules", response_model=List[IPRuleResponse])
async def list_ip_rules(
    enabled_only: bool = True,
    current_user: User = Depends(require_admin)
):
    """List all IP control rules."""
    return await ip_control_service.list_rules(enabled_only)

@router.post("/api/ip-control/rules", response_model=IPRuleResponse)
async def create_ip_rule(
    rule: IPRuleCreate,
    current_user: User = Depends(require_admin)
):
    """Create a new IP control rule."""
    # Validate rule
    await validate_ip_rule(rule)
    
    # Create rule
    created_rule = await ip_control_service.create_rule(
        rule=rule,
        created_by=current_user.id
    )
    
    # Audit log
    await audit_logger.log(
        event_type="security",
        event_action="ip_rule_created",
        event_result="success",
        severity="info",
        actor_id=current_user.id,
        resource_type="ip_rule",
        resource_id=created_rule.id,
        details={"rule_name": rule.name, "mode": rule.mode}
    )
    
    return created_rule

@router.post("/api/ip-control/test")
async def test_ip_access(
    request: IPTestRequest,
    current_user: User = Depends(require_admin)
):
    """Test if an IP would have access to a path."""
    allowed, rule, reason = await ip_control_service.check_access(
        ip=request.ip_address,
        path=request.path,
        user_id=request.test_user_id,
        user_roles=request.test_roles
    )
    
    return {
        "ip_address": request.ip_address,
        "path": request.path,
        "allowed": allowed,
        "matching_rule": rule,
        "reason": reason,
        "country": await geoip_service.get_country(request.ip_address) if geoip_service else None
    }

@router.post("/api/ip-control/block")
async def block_ip_temporarily(
    request: BlockIPRequest,
    current_user: User = Depends(require_admin)
):
    """Temporarily block an IP address."""
    await ip_control_service.block_ip(
        ip=request.ip_address,
        duration_minutes=request.duration_minutes,
        reason=request.reason,
        blocked_by=current_user.id
    )
    
    return {"status": "blocked", "until": request.blocked_until}

@router.get("/api/ip-control/analytics")
async def get_ip_analytics(
    time_range: str = "24h",
    current_user: User = Depends(require_admin)
):
    """Get IP access analytics."""
    return await ip_control_service.get_analytics(time_range)

# Test cases for IP control system
@pytest.mark.asyncio
async def test_ip_control_scenarios():
    # Test 1: Admin access from localhost
    assert await ip_service.check_access(
        ip="127.0.0.1",
        path="/admin/dashboard",
        user_roles=["admin"]
    ) == (True, "admin_localhost_only", "Allowed by allowlist rule")
    
    # Test 2: Admin access from external IP
    assert await ip_service.check_access(
        ip="203.0.113.5",
        path="/admin/dashboard",
        user_roles=["admin"]
    ) == (False, "admin_localhost_only", "Not in allowlist")
    
    # Test 3: Partner API access
    assert await ip_service.check_access(
        ip="203.0.113.10",
        path="/api/v1/partner/data",
        user_roles=["partner"]
    ) == (True, "partner_api_access", "Allowed by allowlist rule")
    
    # Test 4: Geo-blocking
    assert await ip_service.check_access(
        ip="1.2.3.4",  # Assume this resolves to blocked country
        path="/api/public/data",
        user_roles=["user"]
    ) == (False, "geo_restrictions", "Blocked by blocklist rule")
    
    # Test 5: Temporary block
    await ip_service.block_ip("192.168.1.100", 60, "Suspicious activity")
    assert await ip_service.check_access(
        ip="192.168.1.100",
        path="/api/any",
        user_roles=["user"]
    ) == (False, "temporary_block", "IP is temporarily blocked")