"""
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 jwt
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]] = {}
    
    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):
                # 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
                
                # Execute the original function
                return await func(request, *args, **kwargs)
            
            return wrapper
        return decorator
