""" 安全 middleware for the Pixel API.

Provides:
- Suspicious activity detection
- Failed login tracking
- Automatic IP blocking
- Anomaly detection for request patterns
- Input sanitization for all requests
""" 
import time
import logging
import json
from typing import Callable, Optional, Dict, Set
from collections import defaultdict
from datetime import datetime, timedelta
from fastapi import Request, Response, status
from fastapi.responses import JSONResponse
from redis import asyncio as aioredis
from redis.exceptions import RedisError

from src.utils.errors import ErrorCode, InvalidParameterException
from src.utils.auth_logging import AuthLogger, AuthEventType
from src.utils.validators import sanitize_string, sanitize_dict

logger = logging.getLogger(__name__)


class SecurityMonitor:
    """ 监视器 for suspicious activity and security threats.
    
    Tracks:
    - Failed login attempts
    - Unusual request patterns
    - Rapid requests from same IP
    - Access to sensitive endpoints
    """
    
    def __init__(self, redis_url: str = "redis://localhost:6379"):
        """ Initialize security monitor.
        
        Args:
            redis_url: Redis connection URL
        """
        self.redis_url = redis_url
        self._redis: Optional[aioredis.Redis] = None
        self._connected = False
        
        # 配置
        self.max_failed_logins = 5  # Max failed logins before blocking
        self.failed_login_window = 300  # 时间 window in seconds (5 minutes)
        self.block_duration = 3600  # Block duration in seconds (1 hour)
        
        self.max_requests_per_minute = 6000  # Increased for dev/elasticity testing
        self.suspicious_request_threshold = 10000  # Increased for dev/elasticity testing
        
        # In-memory tracking (fallback if Redis unavailable)
        self._failed_logins: Dict[str, list] = defaultdict(list)
        self._blocked_ips: Set[str] = set()
        self._request_counts: Dict[str, list] = defaultdict(list)
    
    async def connect(self):
        """Establish connection to Redis."""
        if not self._connected:
            try:
                self._redis = await aioredis.from_url(
                    self.redis_url,
                    decode_responses=False
                )
                await self._redis.ping()
                self._connected = True
                logger.info(f"Security monitor connected to Redis at {self.redis_url}")
            except RedisError as e:
                logger.error(f"Failed to connect to Redis for security monitoring: {e}")
                self._connected = False
    
    async def disconnect(self):
        """Close connection to Redis."""
        if self._redis:
            await self._redis.close()
            self._connected = False
            logger.info("Security monitor disconnected from Redis")
    
    async def record_failed_login(self, identifier: str, ip_address: str) -> bool:
        """ 记录 a failed login attempt.
        
        Args:
            identifier: User identifier (username or email)
            ip_address: IP address of the request
            
        Returns:
            True if IP should be blocked, False otherwise
        """
        current_time = time.time()
        key = f"security:failed_login:{identifier}:{ip_address}"
        
        if self._connected:
            try:
                # Add timestamp to sorted set
                await self._redis.zadd(key, {str(current_time): current_time})
                
                # Remove old entries outside the window
                window_start = current_time - self.failed_login_window
                await self._redis.zremrangebyscore(key, 0, window_start)
                
                # 计数 failed attempts in window
                count = await self._redis.zcard(key)
                
                # 集合 expiry on key
                await self._redis.expire(key, self.failed_login_window)
                
                # Check if threshold exceeded
                if count >= self.max_failed_logins:
                    await self._block_ip(ip_address, "excessive_failed_logins")
                    logger.warning(
                        f"IP blocked due to excessive failed logins",
                        extra={
                            'ip_address': ip_address,
                            'identifier': identifier,
                            'failed_attempts': count
                        }
                    )
                    return True
                
                return False
                
            except RedisError as e:
                logger.error(f"Redis error recording failed login: {e}")
                # Fall back to in-memory tracking
                return self._record_failed_login_memory(identifier, ip_address, current_time)
        else:
            # Use in-memory tracking
            return self._record_failed_login_memory(identifier, ip_address, current_time)
    
    def _record_failed_login_memory(self, identifier: str, ip_address: str, current_time: float) -> bool:
        """In-memory fallback for failed login tracking.""" 
        key = f"{identifier}:{ip_address}"
        
        # Add timestamp
        self._failed_logins[key].append(current_time)
        
        # Remove old entries
        window_start = current_time - self.failed_login_window
        self._failed_logins[key] = [
            t for t in self._failed_logins[key]
            if t > window_start
        ]
        
        # Check threshold
        if len(self._failed_logins[key]) >= self.max_failed_logins:
            self._blocked_ips.add(ip_address)
            return True
        
        return False
    
    async def _block_ip(self, ip_address: str, reason: str):
        """
        Block an IP address.
        
        Args:
            ip_address: IP address to block
            reason: Reason for blocking
        """ 
        key = f"security:blocked_ip:{ip_address}"
        
        if self._connected:
            try:
                await self._redis.setex(
                    key,
                    self.block_duration,
                    reason
                )
            except RedisError as e:
                logger.error(f"Redis error blocking IP: {e}")
                self._blocked_ips.add(ip_address)
        else:
            self._blocked_ips.add(ip_address)
    
    async def is_ip_blocked(self, ip_address: str) -> tuple[bool, Optional[str]]:
        """
        Check if an IP address is blocked.
        
        Args:
            ip_address: IP address to check
            
        Returns:
            Tuple of (is_blocked, reason)
        """ 
        key = f"security:blocked_ip:{ip_address}"
        
        if self._connected:
            try:
                reason = await self._redis.get(key)
                if reason:
                    return True, reason.decode() if isinstance(reason, bytes) else reason
                return False, None
            except RedisError as e:
                logger.error(f"Redis error checking blocked IP: {e}")
                # Fall back to in-memory
                return ip_address in self._blocked_ips, "blocked" if ip_address in self._blocked_ips else None
        else:
            return ip_address in self._blocked_ips, "blocked" if ip_address in self._blocked_ips else None
    
    async def record_request(self, ip_address: str) -> bool:
        """ 记录 a request and check for suspicious patterns.
        
        Args:
            ip_address: IP address of the request
            
        Returns:
            True if activity is suspicious, False otherwise
        """
        current_time = time.time()
        key = f"security:requests:{ip_address}"
        
        if self._connected:
            try:
                # Add timestamp to sorted set
                await self._redis.zadd(key, {str(current_time): current_time})
                
                # Remove old entries (older than 1 minute)
                window_start = current_time - 60
                await self._redis.zremrangebyscore(key, 0, window_start)
                
                # 计数 requests in last minute
                count = await self._redis.zcard(key)
                
                # 集合 expiry
                await self._redis.expire(key, 60)
                
                # Check for suspicious activity
                if count > self.suspicious_request_threshold:
                    logger.warning(
                        f"Suspicious activity detected",
                        extra={
                            'ip_address': ip_address,
                            'requests_per_minute': count,
                            'threshold': self.suspicious_request_threshold
                        }
                    )
                    await self._block_ip(ip_address, "suspicious_request_pattern")
                    return True
                
                return False
                
            except RedisError as e:
                logger.error(f"Redis error recording request: {e}")
                return False
        else:
            # In-memory tracking
            self._request_counts[ip_address].append(current_time)
            
            # Remove old entries
            window_start = current_time - 60
            self._request_counts[ip_address] = [
                t for t in self._request_counts[ip_address]
                if t > window_start
            ]
            
            # Check threshold
            if len(self._request_counts[ip_address]) > self.suspicious_request_threshold:
                self._blocked_ips.add(ip_address)
                return True
            
            return False
    
    async def unblock_ip(self, ip_address: str):
        """
        Manually unblock an IP address.
        
        Args:
            ip_address: IP address to unblock
        """ 
        key = f"security:blocked_ip:{ip_address}"
        
        if self._connected:
            try:
                await self._redis.delete(key)
                logger.info(f"IP unblocked: {ip_address}")
            except RedisError as e:
                logger.error(f"Redis error unblocking IP: {e}")
        
        # Also remove from in-memory set
        self._blocked_ips.discard(ip_address)


# 全局的 security monitor instance
_security_monitor: Optional[SecurityMonitor] = None


def get_security_monitor() -> SecurityMonitor:
    """ Get the global security monitor instance.
    
    Returns:
        SecurityMonitor instance
    """ 
    global _security_monitor
    if _security_monitor is None:
        from src.config.settings import REDIS_URL
        _security_monitor = SecurityMonitor(redis_url=REDIS_URL)
    return _security_monitor


async def security_middleware(request: Request, call_next: Callable) -> Response:
    """ 安全 middleware for detecting and blocking suspicious activity.
    
    Checks:
    - IP blocking status
    - Request patterns
    - Input sanitization
    """ 
    monitor = get_security_monitor()
    
    # Get client IP
    forwarded_for = request.headers.get("X-Forwarded-For")
    if forwarded_for:
        ip_address = forwarded_for.split(",")[0].strip()
    else:
        ip_address = request.client.host if request.client else "unknown"
    
    # Check if IP is blocked
    is_blocked, reason = await monitor.is_ip_blocked(ip_address)
    if is_blocked:
        logger.warning(
            f"Blocked IP attempted access",
            extra={
                'ip_address': ip_address,
                'path': request.url.path,
                'reason': reason
            }
        )
        
        return JSONResponse(
            status_code=status.HTTP_403_FORBIDDEN,
            content={
                "error_code": ErrorCode.AUTHORIZATION_ERROR,
                "message": "Access denied due to suspicious activity",
                "details": {
                    "reason": "Your IP has been temporarily blocked"
                }
            }
        )
    
    # 清理 request body for POST/PUT/PATCH requests
    if request.method in ["POST", "PUT", "PATCH"]:
        try:
            # Read the body
            body = await request.body()
            if body:
                try:
                    # 解析 JSON body
                    data = json.loads(body)
                    
                    # 清理 the data
                    sanitized_data = sanitize_dict(data, allow_html=False)
                    
                    # Replace the body with sanitized data
                    # 注意: We need to create a new request with sanitized body
                    # FastAPI doesn't allow modifying request body directly
                    # So we'll store sanitized data in request.state for controllers to use
                    request.state.sanitized_body = sanitized_data
                    
                except json.JSONDecodeError:
                    # Not JSON, skip sanitization
                    pass
                except InvalidParameterException as e:
                    # Malicious input detected
                    logger.warning(
                        f"Malicious input detected",
                        extra={
                            'ip_address': ip_address,
                            'path': request.url.path,
                            'error': str(e)
                        }
                    )
                    
                    return JSONResponse(
                        status_code=status.HTTP_400_BAD_REQUEST,
                        content={
                            "error_code": ErrorCode.INVALID_PARAMETER,
                            "message": "Invalid input detected",
                            "details": {
                                "error": str(e)
                            }
                        }
                    )
        except Exception as e:
            logger.error(f"Error sanitizing request body: {e}")
    
    # 记录 request for pattern analysis
    await monitor.record_request(ip_address)
    
    # 进程 request
    response = await call_next(request)
    
    return response


def setup_security_middleware(app):
    """ Setup security middleware for the FastAPI application.
    
    Args:
        app: FastAPI application instance
    """
    app.middleware("http")(security_middleware)
    logger.info("Security middleware registered")


async def init_security_monitor():
    """ Initialize security monitor connection on startup."""
    monitor = get_security_monitor()
    await monitor.connect()


async def shutdown_security_monitor():
    """ 清理up security monitor connection on shutdown."""
    monitor = get_security_monitor()
    await monitor.disconnect()


async def security_headers_middleware(request: Request, call_next: Callable) -> Response:
    """
    Add security headers to all responses.
    
    Headers added:
    - Strict-Transport-Security (HSTS)
    - Content-Security-Policy (CSP)
    - X-Frame-Options
    - X-Content-Type-Options
    - X-XSS-Protection
    - Referrer-Policy
    """ 
    response = await call_next(request)
    
    # HSTS - Force HTTPS for 1 year
    response.headers["Strict-Transport-Security"] = "max-age=31536000; includeSubDomains"
    
    # CSP - Restrict resource loading
    # 注意: Adjust this based on your frontend requirements
    csp_directives = [
        "default-src 'self'",
        "script-src 'self' 'unsafe-inline' 'unsafe-eval'",  # Allow inline scripts for development
        "style-src 'self' 'unsafe-inline'",  # Allow inline styles
        "img-src 'self' data: https:",  # Allow images from self, data URIs, and HTTPS
        "font-src 'self' data:",
        "connect-src 'self' https:",  # Allow API calls to self and HTTPS endpoints
        "frame-ancestors 'none'",  # Prevent framing (same as X-Frame-Options)
        "base-uri 'self'",
        "form-action 'self'",
    ]
    response.headers["Content-Security-Policy"] = "; ".join(csp_directives)
    
    # X-Frame-Options - Prevent clickjacking
    response.headers["X-Frame-Options"] = "DENY"
    
    # X-Content-Type-Options - Prevent MIME sniffing
    response.headers["X-Content-Type-Options"] = "nosniff"
    
    # X-XSS-Protection - Enable XSS filter
    response.headers["X-XSS-Protection"] = "1; mode=block"
    
    # Referrer-Policy - Control referrer information
    response.headers["Referrer-Policy"] = "strict-origin-when-cross-origin"
    
    # Permissions-Policy - Control browser features
    response.headers["Permissions-Policy"] = "geolocation=(), microphone=(), camera=()"
    
    return response


def setup_security_headers_middleware(app):
    """ Setup security headers middleware for the FastAPI application.
    
    Args:
        app: FastAPI application instance
    """
    app.middleware("http")(security_headers_middleware)
    logger.info("Security headers middleware registered")
