"""
FastAPI middleware for payload-guard.
Provides request/response filtering for FastAPI applications.
"""

from __future__ import annotations
import json
from typing import Any, Callable, Optional
from fastapi import FastAPI, Request, Response
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware

from ..core.filter import build_shape, GuardConfig
from ..core.types import ShapeDescriptor, GuardMiddlewareOptions
from ..core.compiler import CompileOptions, compile_shape
from ..core.security import DEFAULT_SENSITIVE_FIELDS


class PayloadGuardMiddleware(BaseHTTPMiddleware):
    """
    FastAPI middleware for payload filtering and sanitization.

    Example:
        from fastapi import FastAPI
        from payloadguard.fastapi import PayloadGuardMiddleware

        app = FastAPI()

        request_shape = {
            'name': 'string',
            'email': 'string',
        }

        app.add_middleware(
            PayloadGuardMiddleware,
            request_shape=request_shape,
            dev_mode=True,
        )
    """

    def __init__(
        self,
        app: FastAPI,
        request_shape: Optional[ShapeDescriptor] = None,
        response_shape: Optional[ShapeDescriptor] = None,
        sensitive_fields: Optional[list[str]] = None,
        dev_mode: bool = False,
        strict: bool = False,
        max_array_length: Optional[int] = None,
        collect_errors: bool = False,
        ignore_routes: Optional[list[str]] = None,
        sanitize_body: bool = True,
        filter_response: bool = False,
        **kwargs,
    ):
        super().__init__(app)
        self.request_shape = request_shape
        self.response_shape = response_shape
        self.sensitive_fields = sensitive_fields or []
        self.dev_mode = dev_mode
        self.strict = strict
        self.max_array_length = max_array_length
        self.collect_errors = collect_errors
        self.ignore_routes = ignore_routes or []
        self.sanitize_body = sanitize_body
        self.filter_response = filter_response
        self.extra_config = kwargs

        # Compile shapes
        self._request_filter = None
        self._response_filter = None

        if request_shape:
            self._request_filter = compile_shape(
                request_shape,
                CompileOptions(
                    sensitive=self.sensitive_fields + DEFAULT_SENSITIVE_FIELDS,
                    dev=dev_mode,
                    strict=strict,
                    max_array_length=max_array_length,
                    collect_errors=collect_errors,
                ),
            )

        if response_shape:
            self._response_filter = compile_shape(
                response_shape,
                CompileOptions(
                    sensitive=self.sensitive_fields + DEFAULT_SENSITIVE_FIELDS,
                    dev=dev_mode,
                    strict=strict,
                    max_array_length=max_array_length,
                    collect_errors=collect_errors,
                ),
            )

    def _should_ignore(self, path: str) -> bool:
        """Check if route should be ignored."""
        for prefix in self.ignore_routes:
            if path.startswith(prefix):
                return True
        return False

    async def dispatch(self, request: Request, call_next: Callable) -> Response:
        """Process request and response."""
        # Check if route should be ignored
        if self._should_ignore(request.url.path):
            return await call_next(request)

        # Sanitize request body
        if self.sanitize_body and self._request_filter and request.method in ("POST", "PUT", "PATCH"):
            try:
                body = await request.json()
                sanitized_body = self._request_filter(body)

                # Create new request with sanitized body
                async def receive() -> dict:
                    return {
                        "type": "http.request",
                        "body": json.dumps(sanitized_body).encode(),
                    }

                request._receive = receive
                request.state.sanitized_body = sanitized_body
            except Exception as e:
                if self.dev_mode:
                    print(f"[payload-guard] Request sanitization error: {e}")

        # Process response
        response = await call_next(request)

        # Filter response
        if self.filter_response and self._response_filter and isinstance(response, JSONResponse):
            try:
                # Get response body
                body = response.body.decode()
                data = json.loads(body)

                # Filter data
                filtered_data = self._response_filter(data)

                # Create new response
                return JSONResponse(
                    content=filtered_data,
                    status_code=response.status_code,
                    headers=dict(response.headers),
                )
            except Exception as e:
                if self.dev_mode:
                    print(f"[payload-guard] Response filtering error: {e}")

        return response


def guard_middleware(
    app: FastAPI,
    request_shape: Optional[ShapeDescriptor] = None,
    response_shape: Optional[ShapeDescriptor] = None,
    options: Optional[GuardMiddlewareOptions] = None,
) -> None:
    """
    Add payload guard middleware to FastAPI app.

    Args:
        app: FastAPI application
        request_shape: Shape for request body sanitization
        response_shape: Shape for response filtering
        options: Middleware options
    """
    if options is None:
        options = GuardMiddlewareOptions()

    app.add_middleware(
        PayloadGuardMiddleware,
        request_shape=request_shape or options.request_shape,
        response_shape=response_shape or options.response_shape,
        sensitive_fields=options.sensitive_fields,
        dev_mode=options.dev_mode or False,
        strict=options.strict or False,
        max_array_length=options.max_array_length,
        collect_errors=options.collect_errors or False,
        ignore_routes=options.ignore_routes if hasattr(options, 'ignore_routes') else None,
        sanitize_body=options.sanitize_body,
        filter_response=options.filter_response,
    )


def route_guard(
    shape: ShapeDescriptor,
    sensitive_fields: Optional[list[str]] = None,
    dev_mode: bool = False,
    strict: bool = False,
) -> Callable:
    """
    Decorator for route-level payload validation.

    Example:
        @app.post("/users")
        @route_guard(user_shape)
        async def create_user(user: dict):
            return user
    """

    def decorator(func: Callable) -> Callable:
        async def wrapper(*args, **kwargs):
            # Compile shape
            request_filter = compile_shape(
                shape,
                CompileOptions(
                    sensitive=(sensitive_fields or []) + DEFAULT_SENSITIVE_FIELDS,
                    dev=dev_mode,
                    strict=strict,
                ),
            )

            # Get request from kwargs or args
            request = kwargs.get("request")
            if request:
                try:
                    body = await request.json()
                    sanitized = request_filter(body)
                    kwargs["sanitized_body"] = sanitized
                except Exception as e:
                    if dev_mode:
                        print(f"[payload-guard] Route guard error: {e}")

            return await func(*args, **kwargs)

        return wrapper

    return decorator
