"""Starlette integration for LeWAF."""

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING, Any, cast
from urllib.parse import parse_qs

from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import JSONResponse, RedirectResponse, Response

from lewaf.config.loader import ConfigLoader
from lewaf.integration import WAF

if TYPE_CHECKING:
    from collections.abc import Awaitable, Callable

    from starlette.applications import Starlette
    from starlette.requests import Request

logger = logging.getLogger(__name__)


class LeWAFMiddleware(BaseHTTPMiddleware):
    """Starlette middleware for LeWAF protection."""

    def __init__(
        self,
        app: Starlette,
        waf: WAF | None = None,
        config_file: str | None = None,
        rules: list[str] | None = None,
        block_response_status: int = 403,
        block_response_body: str = "Request blocked by WAF",
    ):
        super().__init__(app)

        if waf is not None:
            self.waf = waf
        else:
            # Create WAF from config
            config: dict[str, Any] = {}
            if config_file:
                # Load rules from config file
                loader = ConfigLoader()
                loaded_config = loader.load_from_file(Path(config_file))
                config["rules"] = loaded_config.rules
                config["rule_files"] = loaded_config.rule_files
            elif rules:
                config["rules"] = rules
            else:
                # Default basic rules
                config["rules"] = [
                    'SecRule ARGS:test "@rx attack" "id:1001,phase:1,deny,log,msg:\'Test attack detected\'"',
                    'SecRule REQUEST_HEADERS:User-Agent "@rx (bot|scanner|crawler)" "id:1002,phase:1,deny,log,msg:\'Bot detected\'"',
                ]

            self.waf = WAF(config)

        self.block_response_status = block_response_status
        self.block_response_body = block_response_body

    async def dispatch(
        self, request: Request, call_next: Callable[[Request], Awaitable[Response]]
    ) -> Response:
        """Process request through WAF before forwarding."""

        # Create transaction
        tx = self.waf.new_transaction()

        try:
            # Process URI and method
            tx.process_uri(str(request.url.path), request.method)

            # Add query parameters to ARGS
            if request.url.query:
                query_params = parse_qs(request.url.query)
                for key, values in query_params.items():
                    for value in values:
                        tx.variables.args.add(key, value)

            # Add request headers
            for name, value in request.headers.items():
                tx.variables.request_headers.add(name, value)

            # Process request headers phase (Phase 1)
            interruption = tx.process_request_headers()
            if interruption:
                logger.warning(
                    f"Request blocked in headers phase by rule {interruption['rule_id']}"
                )
                return self._create_block_response(interruption)

            # Read and process request body (Phase 2)
            body = await request.body()
            if body:
                content_type = request.headers.get("content-type", "")
                tx.add_request_body(body, content_type)

            body_interruption = tx.process_request_body()
            if body_interruption:
                logger.warning(
                    f"Request blocked in body phase by rule {body_interruption['rule_id']}"
                )
                return self._create_block_response(body_interruption)

            # Request passed WAF, continue to upstream
            response = await call_next(request)

            # Process response headers phase (Phase 3)
            response_headers = dict(response.headers)
            tx.add_response_headers(response_headers)
            tx.add_response_status(response.status_code)

            response_interruption = tx.process_response_headers()
            if response_interruption:
                logger.warning(
                    f"Response blocked in headers phase by rule {response_interruption['rule_id']}"
                )
                return self._create_block_response(response_interruption)

            # Note: Response body phase (Phase 4) would require buffering the entire
            # response body, which has performance implications. Skipped for now.

            return response

        except Exception as e:
            logger.error("Error in LeWAF middleware: %s", e)
            # On error, allow request to proceed (fail open)
            return await call_next(request)

    def _create_block_response(self, interruption: dict[str, Any]) -> Response:
        """Create a blocked request response."""
        # Check for redirect URL
        redirect_url = interruption.get("redirect_url")
        if redirect_url:
            return RedirectResponse(url=redirect_url, status_code=302)

        if self.block_response_status == 403:
            return JSONResponse(
                status_code=self.block_response_status,
                content={
                    "error": "Request blocked by WAF",
                    "rule_id": interruption["rule_id"],
                    "message": self.block_response_body,
                },
            )
        return Response(
            content=self.block_response_body,
            status_code=self.block_response_status,
            media_type="text/plain",
        )


def create_waf_app(
    target_app: Starlette,
    config_file: str | None = None,
    rules: list[str] | None = None,
    **middleware_kwargs: Any,
) -> Starlette:
    """Create a Starlette app with WAF protection."""

    # Add LeWAF middleware to the target app
    target_app.add_middleware(
        cast("Any", LeWAFMiddleware),
        config_file=config_file,
        rules=rules,
        **middleware_kwargs,
    )

    return target_app


# Backward compatibility alias
CorazaMiddleware = LeWAFMiddleware
