"""
PromptShield — FastAPI / Starlette ASGI Middleware
===================================================

Maintainer : Joseph Mutua
Version    : 1.0.0

Overview
--------
Drop-in ASGI middleware that intercepts every POST/PUT/PATCH request,
scans JSON body fields with PromptShield, and returns HTTP 400 before
the route handler ever runs — if an injection is detected.

Your route handlers require zero changes. On clean requests they see
the original body exactly as sent. On blocked requests they are never
called at all.

Usage
-----
    from fastapi import FastAPI
    from promptshield import PromptShield, Policy
    from promptshield.fastapi_middleware import PromptShieldMiddleware

    app    = FastAPI()
    shield = PromptShield(policy=Policy.STRICT)

    app.add_middleware(
        PromptShieldMiddleware,
        shield        = shield,
        scan_fields   = ["message", "prompt", "input"],
        exclude_paths = ["/health", "/docs", "/openapi.json"],
    )

    @app.post("/chat")
    async def chat(request: ChatRequest):
        # Reached only if request.message passed PromptShield.
        # Unsafe messages are rejected at the middleware layer (HTTP 400).
        response = await call_llm(request.message)
        return {"response": response}

HTTP Response on Block
-----------------------
    HTTP 400 Bad Request
    Content-Type: application/json

    {
        "error":        "Request blocked by PromptShield",
        "threat_level": "critical",
        "score":        0.855,
        "reason":       "Critical injection attack blocked — Role Hijacking."
    }

    Custom headers:
        X-PromptShield: blocked
        X-PromptShield-Score: 0.855
        X-PromptShield-Level: critical

How body buffering works
------------------------
ASGI delivers request bodies as an async stream of chunks. This middleware:
  1. Drains and buffers the full body before the route handler runs.
  2. Inspects the JSON fields.
  3. On block: sends HTTP 400 immediately — the route handler is never called.
  4. On pass:  replays the buffered messages to the underlying app exactly
               as received. The route handler sees the original body.
"""

from __future__ import annotations

import json
import logging
from typing import List, Optional

logger = logging.getLogger(__name__)


class PromptShieldMiddleware:
    """
    Starlette / FastAPI ASGI middleware for PromptShield.

    Protects all POST, PUT, and PATCH routes that receive a JSON body.
    GET, HEAD, OPTIONS, and DELETE requests are passed through untouched.

    Args:
        app           : The ASGI application to wrap.
        shield        : An initialised PromptShield instance.
        scan_fields   : JSON body field names to scan on each request.
                        Default: ["message","prompt","input","query",
                                  "text","content"]
        exclude_paths : URL paths to skip entirely (health checks, docs,
                        static assets).
                        Default: ["/health","/docs","/openapi.json",
                                  "/redoc","/favicon.ico"]

    Example:
        app.add_middleware(
            PromptShieldMiddleware,
            shield        = PromptShield(policy=Policy.STRICT),
            scan_fields   = ["message", "prompt"],
            exclude_paths = ["/health", "/metrics"],
        )
    """

    def __init__(
        self,
        app,
        shield,
        scan_fields:   Optional[List[str]] = None,
        exclude_paths: Optional[List[str]] = None,
    ) -> None:
        self.app           = app
        self.shield        = shield
        self.scan_fields   = scan_fields or [
            "message", "prompt", "input", "query", "text", "content",
        ]
        self.exclude_paths = set(exclude_paths or [
            "/health", "/docs", "/openapi.json", "/redoc", "/favicon.ico",
        ])

        logger.info(
            "PromptShieldMiddleware registered | "
            "scan_fields=%s | excluded_paths=%d",
            self.scan_fields,
            len(self.exclude_paths),
        )

    async def __call__(self, scope, receive, send) -> None:
        """
        ASGI entry point called for every connection.

        Only HTTP requests with mutation methods are inspected.
        Everything else is forwarded to the underlying app unchanged.
        """
        # Pass through non-HTTP connections (WebSocket, lifespan, etc.)
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return

        method = scope.get("method", "GET")
        path   = scope.get("path", "")

        # Only inspect state-changing requests
        if method not in ("POST", "PUT", "PATCH"):
            await self.app(scope, receive, send)
            return

        # Skip explicitly excluded paths
        if path in self.exclude_paths:
            logger.debug(
                "PromptShieldMiddleware: skipping excluded path '%s'", path
            )
            await self.app(scope, receive, send)
            return

        # ------------------------------------------------------------------
        # Buffer the request body
        # ------------------------------------------------------------------
        # We must read the full body before inspecting it, then replay the
        # stream to the route handler if the request is clean.
        body_chunks:       List[bytes] = []
        buffered_messages: List[dict]  = []

        while True:
            message = await receive()
            buffered_messages.append(message)

            if message["type"] == "http.request":
                body_chunks.append(message.get("body", b""))
                if not message.get("more_body", False):
                    break  # All chunks received

        body_bytes = b"".join(body_chunks)

        # ------------------------------------------------------------------
        # Inspect JSON fields
        # ------------------------------------------------------------------
        blocked        = False
        block_response = {}
        block_score    = 0.0
        block_level    = ""

        # Determine Content-Type from request headers
        content_type = ""
        for header_name, header_value in scope.get("headers", []):
            if header_name == b"content-type":
                content_type = header_value.decode("utf-8", errors="ignore").lower()
                break

        if "application/json" in content_type:
            try:
                data = json.loads(body_bytes.decode("utf-8", errors="ignore"))

                if isinstance(data, dict):
                    for field_name in self.scan_fields:
                        value = data.get(field_name)
                        if not isinstance(value, str):
                            continue  # Skip non-string fields

                        result = self.shield.scan(value)

                        if not result.allowed:
                            blocked        = True
                            block_score    = result.score
                            block_level    = result.threat_level
                            block_response = {
                                "error":        "Request blocked by PromptShield",
                                "threat_level": result.threat_level,
                                "score":        result.score,
                                "reason":       result.summary,
                            }

                            logger.warning(
                                "PromptShieldMiddleware [BLOCKED] "
                                "path=%s field=%s score=%.3f level=%s",
                                path, field_name,
                                result.score, result.threat_level,
                            )

                            # Fire on_block callback if configured
                            if self.shield.on_block:
                                try:
                                    self.shield.on_block(result)
                                except Exception as cb_exc:
                                    logger.error(
                                        "on_block callback error: %s",
                                        cb_exc, exc_info=True,
                                    )

                            break  # Stop at first blocked field

            except (json.JSONDecodeError, UnicodeDecodeError) as exc:
                # Non-JSON body — pass through unchanged
                logger.debug(
                    "PromptShieldMiddleware: non-JSON body (%s) — "
                    "passing through unchanged",
                    exc,
                )

        # ------------------------------------------------------------------
        # Return 400 if blocked; replay body to app if clean
        # ------------------------------------------------------------------
        if blocked:
            response_body = json.dumps(block_response).encode("utf-8")
            await send({
                "type":    "http.response.start",
                "status":  400,
                "headers": [
                    [b"content-type",         b"application/json"],
                    [b"content-length",       str(len(response_body)).encode()],
                    [b"x-promptshield",       b"blocked"],
                    [b"x-promptshield-score", str(block_score).encode()],
                    [b"x-promptshield-level", block_level.encode()],
                ],
            })
            await send({"type": "http.response.body", "body": response_body})
            return

        # Clean request — replay the buffered ASGI messages to the route handler
        logger.debug(
            "PromptShieldMiddleware [PASS] path=%s — forwarding to app", path
        )
        msg_iter = iter(buffered_messages)

        async def _replay_receive():
            try:
                return next(msg_iter)
            except StopIteration:
                # Stream ended — return disconnect message
                return {"type": "http.disconnect"}

        await self.app(scope, _replay_receive, send)
