from fastapi import Request, HTTPException, status
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
from typing import Optional
import secrets
import hmac
import hashlib
from ...shared.config.logging_config import get_logger
from ...shared.config.settings import settings

logger = get_logger(__name__)

# Safe HTTP methods that don't require CSRF protection
SAFE_METHODS = {'GET', 'HEAD', 'OPTIONS'}


class CSRFProtectionMiddleware(BaseHTTPMiddleware):
    """
    CSRF Protection Middleware
    
    Validates CSRF tokens for state-changing requests (POST, PUT, DELETE, PATCH).
    Uses Double Submit Cookie pattern for stateless CSRF protection.
    """
    
    CSRF_TOKEN_COOKIE_NAME = 'XSRF-TOKEN'
    CSRF_TOKEN_HEADER_NAME = 'X-XSRF-TOKEN'
    CSRF_SECRET_LENGTH = 32
    
    async def dispatch(self, request: Request, call_next):
        path = request.url.path
        
        # Skip CSRF protection for certain endpoints that don't need it
        # (e.g., public APIs, webhooks with their own validation)
        is_exempt = self._is_exempt_path(path)
        
        # Get or generate CSRF token (always generate for all requests to ensure cookie is set)
        csrf_token = request.cookies.get(self.CSRF_TOKEN_COOKIE_NAME)
        if not csrf_token:
            csrf_token = self._generate_token()
        
        # Skip CSRF validation for safe methods (GET, HEAD, OPTIONS) and exempt paths
        if request.method in SAFE_METHODS or is_exempt:
            response = await call_next(request)
        else:
            # For state-changing requests, validate the token
            if request.method in {'POST', 'PUT', 'DELETE', 'PATCH'}:
                header_token = request.headers.get(self.CSRF_TOKEN_HEADER_NAME)
                
                if not header_token:
                    logger.warning(f"CSRF token missing in header for {request.method} {path}")
                    # Create error response with CSRF cookie set so frontend can retry
                    from fastapi.responses import JSONResponse
                    error_response = JSONResponse(
                        status_code=status.HTTP_403_FORBIDDEN,
                        content={"status": "error", "message": "CSRF token missing. Please include X-XSRF-TOKEN header."}
                    )
                    # Set cookie even on error so client can get the token and retry
                    if not request.cookies.get(self.CSRF_TOKEN_COOKIE_NAME):
                        error_response.set_cookie(
                            key=self.CSRF_TOKEN_COOKIE_NAME,
                            value=csrf_token,
                            httponly=False,
                            secure=settings.is_production,
                            samesite='lax',  # Changed to 'lax' for better cross-origin support
                            max_age=86400 * 7,
                            path='/'
                        )
                    return error_response
                
                # Validate token using constant-time comparison
                if not self._verify_token(csrf_token, header_token):
                    logger.warning(f"CSRF token validation failed for {request.method} {path}")
                    # Create error response with CSRF cookie set so frontend can retry
                    from fastapi.responses import JSONResponse
                    error_response = JSONResponse(
                        status_code=status.HTTP_403_FORBIDDEN,
                        content={"status": "error", "message": "Invalid CSRF token. Please refresh the page and try again."}
                    )
                    # Set cookie even on error so client can get the token and retry
                    if not request.cookies.get(self.CSRF_TOKEN_COOKIE_NAME):
                        error_response.set_cookie(
                            key=self.CSRF_TOKEN_COOKIE_NAME,
                            value=csrf_token,
                            httponly=False,
                            secure=settings.is_production,
                            samesite='lax',  # Changed to 'lax' for better cross-origin support
                            max_age=86400 * 7,
                            path='/'
                        )
                    return error_response
            
            # Process request
            response = await call_next(request)
        
        # Always set CSRF token cookie if not present (ensures client always has it)
        # This allows frontend to read it from cookies for subsequent requests
        if not request.cookies.get(self.CSRF_TOKEN_COOKIE_NAME):
            # Set secure cookie with SameSite protection
            response.set_cookie(
                key=self.CSRF_TOKEN_COOKIE_NAME,
                value=csrf_token,
                httponly=False,  # Must be accessible to JavaScript for header submission
                secure=settings.is_production,  # HTTPS only in production
                samesite='lax',  # Changed to 'lax' for better cross-origin support
                max_age=86400 * 7,  # 7 days
                path='/'
            )
        
        return response
    
    def _is_exempt_path(self, path: str) -> bool:
        """
        Check if path is exempt from CSRF protection.
        
        Exempt paths:
        - Authentication endpoints (login, register, logout, refresh token)
        - Webhook endpoints (they have their own signature validation)
        - Health check endpoints
        - Static file endpoints
        """
        exempt_patterns = [
            '/api/auth/',  # All authentication endpoints
            '/api/webhooks/',
            '/api/stripe/webhook',
            '/api/payments/stripe/webhook',
            '/api/paypal/webhook',
            '/health',
            '/api/health',
            '/static/',
            '/docs',
            '/redoc',
            '/openapi.json'
        ]
        
        return any(path.startswith(pattern) for pattern in exempt_patterns)
    
    def _generate_token(self) -> str:
        """Generate a secure random CSRF token."""
        return secrets.token_urlsafe(self.CSRF_SECRET_LENGTH)
    
    def _verify_token(self, cookie_token: str, header_token: str) -> bool:
        """
        Verify CSRF token using constant-time comparison.
        
        Uses the Double Submit Cookie pattern - the token in the cookie
        must match the token in the header.
        """
        if not cookie_token or not header_token:
            return False
        
        # Constant-time comparison to prevent timing attacks
        return hmac.compare_digest(cookie_token, header_token)


def get_csrf_token(request: Request) -> Optional[str]:
    """Helper function to get CSRF token from request cookies."""
    return request.cookies.get(CSRFProtectionMiddleware.CSRF_TOKEN_COOKIE_NAME)

