# -*- coding: utf-8 -*-
"""HTTP Authentication Middleware.

This middleware allows plugins to:
1. Transform request headers before authentication (HTTP_PRE_REQUEST)
2. Inspect responses after request completion (HTTP_POST_REQUEST)
"""

# Standard
import logging

# Third-Party
from fastapi import Request
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp

# First-Party
from mcpgateway.plugins.framework import GlobalContext, HttpHeaderPayload, HttpHookType, HttpPostRequestPayload, HttpPreRequestPayload, PluginManager
from mcpgateway.utils.correlation_id import generate_correlation_id, get_correlation_id

logger = logging.getLogger(__name__)


class HttpAuthMiddleware(BaseHTTPMiddleware):
    """Middleware for HTTP authentication hooks.

    This middleware invokes plugin hooks for HTTP request processing:
    - HTTP_PRE_REQUEST: Before any authentication, allows header transformation
    - HTTP_POST_REQUEST: After request completion, allows response inspection

    The middleware allows plugins to:
    - Convert custom authentication tokens to standard formats
    - Add tracing/correlation headers
    - Implement custom authentication schemes
    - Audit authentication attempts
    - Log response status and headers
    """

    def __init__(self, app: ASGIApp, plugin_manager: PluginManager | None = None):
        """Initialize the HTTP auth middleware.

        Args:
            app: The ASGI application
            plugin_manager: Optional plugin manager for hook invocation
        """
        super().__init__(app)
        self.plugin_manager = plugin_manager

    async def dispatch(self, request: Request, call_next):
        """Process request through plugin hooks.

        Args:
            request: The incoming request
            call_next: The next middleware/handler in the chain

        Returns:
            The response from the application
        """
        # Skip hook invocation if no plugin manager
        if not self.plugin_manager:
            return await call_next(request)

        # Skip payload creation if no HTTP hooks registered
        has_pre = self.plugin_manager.has_hooks_for(HttpHookType.HTTP_PRE_REQUEST)
        has_post = self.plugin_manager.has_hooks_for(HttpHookType.HTTP_POST_REQUEST)

        if not has_pre and not has_post:
            return await call_next(request)

        # Use correlation ID from CorrelationIDMiddleware if available
        # This ensures all hooks and downstream code see the same unified request ID
        request_id = get_correlation_id()
        if not request_id:
            # Fallback if correlation ID middleware is disabled
            request_id = generate_correlation_id()
            logger.debug(f"Correlation ID not found, generated fallback: {request_id}")

        request.state.request_id = request_id

        # Create global context for hooks
        global_context = GlobalContext(
            request_id=request_id,
            server_id=None,  # Not specific to any server
            tenant_id=None,  # Not specific to any tenant
        )

        # Extract client information
        client_host = None
        client_port = None
        if request.client:
            client_host = request.client.host
            client_port = request.client.port

        # Initialize context_table for potential use by POST hook
        context_table = None

        # PRE-REQUEST HOOK: Allow plugins to transform headers before authentication
        # Only create payload and invoke hook if plugins are registered for this hook type
        if has_pre:
            try:
                pre_result, context_table = await self.plugin_manager.invoke_hook(
                    HttpHookType.HTTP_PRE_REQUEST,
                    payload=HttpPreRequestPayload(
                        path=str(request.url.path),
                        method=request.method,
                        headers=HttpHeaderPayload(root=dict(request.headers)),
                        client_host=client_host,
                        client_port=client_port,
                    ),
                    global_context=global_context,
                    local_contexts=None,
                    violations_as_exceptions=False,  # Don't block on pre-request violations
                )

                if context_table:
                    request.state.plugin_context_table = context_table

                if global_context:
                    request.state.plugin_global_context = global_context

                # Apply modified headers if plugin returned them
                if pre_result.modified_payload:
                    # Modify request headers by updating request.scope["headers"]
                    # This is the proper way to modify headers in Starlette/FastAPI
                    # Reference: https://stackoverflow.com/questions/69934160/python-how-to-manipulate-fastapi-request-headers-to-be-mutable
                    modified_headers_dict = pre_result.modified_payload.root

                    # Merge modified headers with original headers (modified headers take precedence)
                    original_headers = dict(request.headers)
                    merged_headers = {**original_headers, **modified_headers_dict}

                    # Update request.scope["headers"] which is the raw header list Starlette uses
                    # Convert dict to list of (name, value) tuples with lowercase byte keys
                    request.scope["headers"] = [(name.lower().encode(), value.encode()) for name, value in merged_headers.items()]

                    logger.debug(f"Pre-request hook modified headers: {list(modified_headers_dict.keys())}")

            except Exception as e:
                # Log but don't fail the request if pre-hook has issues
                logger.warning(f"HTTP_PRE_REQUEST hook failed: {e}", exc_info=True)

        # Process the request through the rest of the application
        response = await call_next(request)

        # POST-REQUEST HOOK: Allow plugins to inspect and modify response
        # Only create payload and invoke hook if plugins are registered for this hook type
        if has_post:
            try:
                # Extract response headers
                response_headers = HttpHeaderPayload(root=dict(response.headers))

                post_result, _ = await self.plugin_manager.invoke_hook(
                    HttpHookType.HTTP_POST_REQUEST,
                    payload=HttpPostRequestPayload(
                        path=str(request.url.path),
                        method=request.method,
                        headers=HttpHeaderPayload(root=dict(request.headers)),
                        client_host=client_host,
                        client_port=client_port,
                        response_headers=response_headers,
                        status_code=response.status_code,
                    ),
                    global_context=global_context,
                    local_contexts=context_table,  # Pass context from pre-hook
                    violations_as_exceptions=False,  # Don't block on post-request violations
                )

                # Apply modified response headers if plugin returned them
                if post_result.modified_payload:
                    modified_response_headers = post_result.modified_payload.root
                    # Update response headers (response.headers is mutable)
                    for header_name, header_value in modified_response_headers.items():
                        response.headers[header_name] = header_value
                    logger.debug(f"Post-request hook modified response headers: {list(modified_response_headers.keys())}")

            except Exception as e:
                # Log but don't fail the response if post-hook has issues
                logger.warning(f"HTTP_POST_REQUEST hook failed: {e}", exc_info=True)

        return response
