"""
Core Guardian class for MCP Guardian security middleware.

This module contains the main Guardian class that provides JWT authentication,
scope-based authorization, and rate limiting for FastAPI applications.
"""

import time
import json
import logging
import jwt
from datetime import datetime, timezone
from functools import wraps
from typing import Dict, List, Any, Optional
from fastapi import HTTPException, Request
from .config import load_config


class Guardian:
    """
    Security middleware for FastAPI applications.
    
    Provides JWT authentication, scope-based authorization, and rate limiting
    through a simple decorator-based API.
    """
    
    def __init__(self, config_path: str):
        """
        Initialize the Guardian with configuration from a YAML file.
        
        Args:
            config_path: Path to the YAML configuration file
        """
        self.config = load_config(config_path)
        self.jwt_config = self.config['jwt']
        self.rate_limit_config = self.config['rate_limit']
        
        # In-memory rate limit tracker: {user_id: [timestamp1, timestamp2, ...]}
        self._rate_limit_tracker: Dict[str, List[float]] = {}
        
        # Initialize dedicated security logger
        self.logger = logging.getLogger('mcp_guardian.security')
        if not self.logger.handlers:
            # Configure logger to output JSON format
            handler = logging.StreamHandler()
            formatter = logging.Formatter('%(message)s')  # Raw JSON output
            handler.setFormatter(formatter)
            self.logger.addHandler(handler)
            self.logger.setLevel(logging.INFO)
            self.logger.propagate = False  # Prevent duplicate logs
    
    def _log_security_event(self, event_action: str, event_status: str, 
                           user_id: Optional[str], source_ip: str, 
                           detail: str, additional_data: Optional[Dict[str, Any]] = None) -> None:
        """
        Log a security event in structured JSON format.
        
        Args:
            event_action: Type of security event (e.g., "authentication_success", "authentication_failed")
            event_status: "SUCCESS" or "FAILURE"
            user_id: User identifier (from JWT sub claim or None if unavailable)
            source_ip: Client IP address
            detail: Human-readable description of the event
            additional_data: Optional additional context data
        """
        log_entry = {
            "timestamp": datetime.now(timezone.utc).isoformat(),
            "event_action": event_action,
            "event_status": event_status,
            "user_id": user_id,
            "source_ip": source_ip,
            "detail": detail,
            "service": "mcp_guardian"
        }
        
        # Add additional context if provided
        if additional_data:
            log_entry.update(additional_data)
        
        # Log as JSON string
        self.logger.info(json.dumps(log_entry))
    
    def _extract_user_id_from_invalid_token(self, token: str) -> Optional[str]:
        """
        Attempt to extract user_id from an invalid token for audit purposes.
        This decodes the token without verification to get the sub claim.
        
        Args:
            token: JWT token string
            
        Returns:
            User ID from sub claim if available, None otherwise
        """
        try:
            # Decode without verification to extract claims
            unverified_payload = jwt.decode(token, options={"verify_signature": False})
            return unverified_payload.get('sub')
        except Exception:
            # If we can't decode even without verification, return None
            return None
    
    def _extract_bearer_token(self, request: Request) -> str:
        """
        Extract Bearer token from Authorization header.
        
        Args:
            request: FastAPI request object
            
        Returns:
            The JWT token string
            
        Raises:
            HTTPException: If token is missing or malformed
        """
        authorization = request.headers.get('Authorization')
        
        if not authorization:
            raise HTTPException(
                status_code=401,
                detail="Missing Authorization header"
            )
        
        if not authorization.startswith('Bearer '):
            raise HTTPException(
                status_code=401,
                detail="Invalid Authorization header format. Expected 'Bearer <token>'"
            )
        
        token = authorization[7:]  # Remove 'Bearer ' prefix
        
        if not token:
            raise HTTPException(
                status_code=401,
                detail="Missing token in Authorization header"
            )
        
        return token
    
    def _validate_jwt(self, token: str) -> Dict[str, Any]:
        """
        Validate JWT token and return the payload.
        
        Args:
            token: JWT token string
            
        Returns:
            Decoded JWT payload
            
        Raises:
            HTTPException: If token validation fails
        """
        try:
            payload = jwt.decode(
                token,
                self.jwt_config['public_key'],
                algorithms=[self.jwt_config['algorithm']],
                audience=self.jwt_config['audience']
            )
            return payload
        except jwt.ExpiredSignatureError:
            raise HTTPException(
                status_code=401,
                detail="Token has expired"
            )
        except jwt.InvalidAudienceError:
            raise HTTPException(
                status_code=401,
                detail="Invalid token audience"
            )
        except jwt.InvalidTokenError as e:
            raise HTTPException(
                status_code=401,
                detail=f"Invalid token: {str(e)}"
            )
    
    def _check_scope(self, payload: Dict[str, Any], required_scope: str) -> None:
        """
        Check if the required scope exists in the JWT payload.
        
        Args:
            payload: Decoded JWT payload
            required_scope: The scope that must be present
            
        Raises:
            HTTPException: If the required scope is not present
        """
        scopes = payload.get('scopes', [])
        
        if not isinstance(scopes, list):
            raise HTTPException(
                status_code=403,
                detail="Invalid scopes claim in token"
            )
        
        if required_scope not in scopes:
            raise HTTPException(
                status_code=403,
                detail=f"Insufficient permissions. Required scope: {required_scope}"
            )
    
    def _check_rate_limit(self, user_id: str) -> None:
        """
        Check if the user has exceeded the rate limit.
        
        Args:
            user_id: User identifier from JWT 'sub' claim
            
        Raises:
            HTTPException: If rate limit is exceeded
        """
        current_time = time.time()
        period_seconds = self.rate_limit_config['period_minutes'] * 60
        max_requests = self.rate_limit_config['requests']
        
        # Clean up old timestamps
        if user_id in self._rate_limit_tracker:
            self._rate_limit_tracker[user_id] = [
                timestamp for timestamp in self._rate_limit_tracker[user_id]
                if current_time - timestamp < period_seconds
            ]
        else:
            self._rate_limit_tracker[user_id] = []
        
        # Check if limit exceeded
        if len(self._rate_limit_tracker[user_id]) >= max_requests:
            raise HTTPException(
                status_code=429,
                detail=f"Rate limit exceeded. Maximum {max_requests} requests per {self.rate_limit_config['period_minutes']} minutes"
            )
        
        # Add current request timestamp
        self._rate_limit_tracker[user_id].append(current_time)
    
    def protect(self, scope: str):
        """
        Decorator to protect FastAPI endpoints with JWT authentication,
        scope-based authorization, and rate limiting.
        
        Args:
            scope: Required scope for accessing the endpoint
            
        Returns:
            Decorated function with security checks
        """
        def decorator(func):
            @wraps(func)
            async def wrapper(request: Request, *args, **kwargs):
                # Get client IP for logging
                client_ip = request.client.host if request.client else "unknown"
                token = None
                user_id = None
                
                try:
                    # Extract and validate Bearer token
                    token = self._extract_bearer_token(request)
                    
                    # Validate JWT
                    payload = self._validate_jwt(token)
                    
                    # Check scope authorization
                    self._check_scope(payload, scope)
                    
                    # Check rate limiting
                    user_id = payload.get('sub')
                    if not user_id:
                        raise HTTPException(
                            status_code=401,
                            detail="Missing 'sub' claim in token"
                        )
                    self._check_rate_limit(user_id)
                    
                    # Attach validated payload to request state
                    request.state.token = payload
                    
                    # Log successful authentication and authorization
                    self._log_security_event(
                        event_action="authentication_success",
                        event_status="SUCCESS",
                        user_id=user_id,
                        source_ip=client_ip,
                        detail=f"User {user_id} successfully authenticated and authorized for scope '{scope}'",
                        additional_data={
                            "scope": scope,
                            "endpoint": f"{request.method} {request.url.path}",
                            "user_agent": request.headers.get("user-agent", "unknown")
                        }
                    )
                    
                    # Execute the original function
                    return await func(request, *args, **kwargs)
                    
                except HTTPException as e:
                    # Determine event action based on status code
                    if e.status_code == 401:
                        event_action = "authentication_failed"
                        # Try to extract user_id from invalid token for audit purposes
                        if token and not user_id:
                            user_id = self._extract_user_id_from_invalid_token(token)
                    elif e.status_code == 403:
                        event_action = "authorization_failed"
                    elif e.status_code == 429:
                        event_action = "rate_limit_exceeded"
                    else:
                        event_action = "security_error"
                    
                    # Log security failure
                    self._log_security_event(
                        event_action=event_action,
                        event_status="FAILURE",
                        user_id=user_id,
                        source_ip=client_ip,
                        detail=f"Security check failed: {e.detail}",
                        additional_data={
                            "status_code": e.status_code,
                            "scope": scope,
                            "endpoint": f"{request.method} {request.url.path}",
                            "user_agent": request.headers.get("user-agent", "unknown"),
                            "error_detail": e.detail
                        }
                    )
                    
                    # Re-raise the exception
                    raise e
            
            return wrapper
        return decorator
