"""PromptGuard middleware -- prompt injection defense for LLM proxy endpoints.

Intercepts POST requests to ``/v1/chat/completions`` and ``/v1/messages``,
extracts message content, and runs a multi-layer detection pipeline:

1. **Rule-based** — regex pattern matching against a curated threat database
2. **Heuristic** — instruction density ratio and entropy analysis

Configurable via ``[security.prompt_guard]`` in ``config.toml`` with four
operating modes: ``off``, ``warn``, ``block``, and ``sanitize``.

This module is part of the TrustShield security platform (Pillar 1: PromptGuard).
"""

from __future__ import annotations

import json
import logging
import math
from collections import Counter
from typing import Any

from pydantic import BaseModel, Field
from starlette.middleware.base import BaseHTTPMiddleware, RequestResponseEndpoint
from starlette.requests import Request
from starlette.responses import JSONResponse, Response

from llmhosts.security.patterns import (
    PATTERNS,
    InjectionPattern,
    compile_custom_patterns,
    detect_base64_injection,
    detect_homoglyph_mixing,
)

logger = logging.getLogger(__name__)

# Endpoints that PromptGuard inspects
_GUARDED_PATHS: frozenset[str] = frozenset({"/v1/chat/completions", "/v1/messages"})


# ---------------------------------------------------------------------------
# Pydantic models
# ---------------------------------------------------------------------------


class ThreatMatch(BaseModel):
    """A single pattern match from the detection pipeline."""

    category: str  # injection, jailbreak, data_exfil, role_hijack, encoding_attack
    pattern_name: str
    confidence: float  # 0.0-1.0
    matched_text: str  # the portion that matched (truncated to 100 chars)


class ThreatAssessment(BaseModel):
    """Aggregated threat assessment for a request."""

    threat_score: float  # 0.0-1.0 (max of individual confidences, boosted if multiple categories)
    categories: list[str]
    details: str
    matches: list[ThreatMatch]
    mode: str  # the configured mode that was applied


class PromptGuardConfig(BaseModel):
    """Configuration for the PromptGuard middleware."""

    enabled: bool = True
    mode: str = "warn"  # off | warn | block | sanitize
    threshold: float = 0.7
    log_all: bool = True
    custom_patterns: list[str] = Field(default_factory=list)


# ---------------------------------------------------------------------------
# Detection engine
# ---------------------------------------------------------------------------

# Instruction-like keywords for density analysis
_INSTRUCTION_WORDS = frozenset(
    {
        "ignore",
        "override",
        "disregard",
        "forget",
        "bypass",
        "disable",
        "pretend",
        "act",
        "role",
        "persona",
        "mode",
        "switch",
        "system",
        "admin",
        "root",
        "sudo",
        "instructions",
        "prompt",
        "reveal",
        "show",
        "print",
        "output",
        "display",
        "dump",
        "new",
        "now",
        "must",
        "always",
        "never",
        "instead",
    }
)

# Threshold for instruction density (ratio of instruction words to total words)
_INSTRUCTION_DENSITY_THRESHOLD = 0.25

# Shannon entropy threshold for detecting encoded/obfuscated content
_ENTROPY_THRESHOLD = 5.0


def _detect_rule_based(text: str, patterns: list[InjectionPattern]) -> list[ThreatMatch]:
    """Run regex-based pattern matching against the text.

    Parameters
    ----------
    text:
        Concatenated message content to scan.
    patterns:
        List of injection patterns to match against.

    Returns
    -------
    list[ThreatMatch]
        All pattern matches found.
    """
    matches: list[ThreatMatch] = []

    for pat in patterns:
        hit = pat.pattern.search(text)
        if hit:
            matched_text = hit.group(0)[:100]
            matches.append(
                ThreatMatch(
                    category=pat.category,
                    pattern_name=pat.name,
                    confidence=pat.confidence,
                    matched_text=matched_text,
                )
            )

    # Base64 encoded instruction detection
    b64_hits = detect_base64_injection(text)
    for decoded, confidence in b64_hits:
        matches.append(
            ThreatMatch(
                category="encoding_attack",
                pattern_name="base64_instruction",
                confidence=confidence,
                matched_text=decoded,
            )
        )

    # Homoglyph mixing detection
    homoglyph_hits = detect_homoglyph_mixing(text)
    for segment, confidence in homoglyph_hits:
        matches.append(
            ThreatMatch(
                category="encoding_attack",
                pattern_name="unicode_homoglyph",
                confidence=confidence,
                matched_text=segment,
            )
        )

    return matches


def _shannon_entropy(text: str) -> float:
    """Calculate Shannon entropy of a text string.

    Higher entropy indicates more randomness, which can signal
    encoded or obfuscated content.
    """
    if not text:
        return 0.0

    counts = Counter(text)
    length = len(text)
    entropy = 0.0

    for count in counts.values():
        probability = count / length
        if probability > 0:
            entropy -= probability * math.log2(probability)

    return entropy


def _instruction_density(text: str) -> float:
    """Calculate the ratio of instruction-like words to total words.

    A high ratio suggests the text is primarily composed of
    instruction/command language rather than natural conversation.
    """
    words = text.lower().split()
    if not words:
        return 0.0

    instruction_count = sum(1 for w in words if w.strip(".,;:!?\"'()") in _INSTRUCTION_WORDS)
    return instruction_count / len(words)


def _detect_heuristic(text: str) -> list[ThreatMatch]:
    """Run heuristic analysis: instruction density ratio and entropy.

    Parameters
    ----------
    text:
        Concatenated message content to scan.

    Returns
    -------
    list[ThreatMatch]
        Heuristic-based threat matches.
    """
    matches: list[ThreatMatch] = []

    # Instruction density check
    density = _instruction_density(text)
    if density >= _INSTRUCTION_DENSITY_THRESHOLD:
        confidence = min(0.5 + density, 0.9)
        matches.append(
            ThreatMatch(
                category="injection",
                pattern_name="high_instruction_density",
                confidence=confidence,
                matched_text=f"instruction_density={density:.2f}",
            )
        )

    # Entropy analysis (for short suspicious blocks)
    # Only flag if text is reasonably short (long natural text has high entropy too)
    if len(text) < 500:
        entropy = _shannon_entropy(text)
        if entropy >= _ENTROPY_THRESHOLD:
            confidence = min(0.4 + (entropy - _ENTROPY_THRESHOLD) * 0.1, 0.75)
            matches.append(
                ThreatMatch(
                    category="encoding_attack",
                    pattern_name="high_entropy",
                    confidence=confidence,
                    matched_text=f"entropy={entropy:.2f}",
                )
            )

    return matches


def _build_assessment(matches: list[ThreatMatch], mode: str) -> ThreatAssessment:
    """Build a ThreatAssessment from a list of matches.

    The threat_score is the maximum individual confidence, boosted by 0.1
    for each additional unique category (capped at 1.0).

    Parameters
    ----------
    matches:
        All matches from rule-based and heuristic detectors.
    mode:
        The configured operating mode.

    Returns
    -------
    ThreatAssessment
        Aggregated assessment.
    """
    if not matches:
        return ThreatAssessment(
            threat_score=0.0,
            categories=[],
            details="No threats detected",
            matches=[],
            mode=mode,
        )

    categories = sorted(set(m.category for m in matches))
    max_confidence = max(m.confidence for m in matches)

    # Boost score for multi-category attacks (0.1 per additional category)
    category_boost = (len(categories) - 1) * 0.1
    threat_score = min(max_confidence + category_boost, 1.0)

    pattern_names = [m.pattern_name for m in matches]
    details = f"Detected {len(matches)} threat(s) in categories: {', '.join(categories)}. Patterns: {', '.join(pattern_names)}"

    return ThreatAssessment(
        threat_score=threat_score,
        categories=categories,
        details=details,
        matches=matches,
        mode=mode,
    )


def _extract_messages(body: dict[str, Any]) -> list[str]:
    """Extract all message content strings from OpenAI or Anthropic format.

    Handles:
    - OpenAI: ``messages[].content`` (string)
    - Anthropic: ``messages[].content`` (string or list of content blocks)

    Parameters
    ----------
    body:
        Parsed JSON request body.

    Returns
    -------
    list[str]
        All extracted message content strings.
    """
    contents: list[str] = []
    messages = body.get("messages")

    if not isinstance(messages, list):
        return contents

    for msg in messages:
        if not isinstance(msg, dict):
            continue

        content = msg.get("content")
        if isinstance(content, str):
            contents.append(content)
        elif isinstance(content, list):
            # Anthropic format: list of content blocks [{type: "text", text: "..."}]
            for block in content:
                if isinstance(block, dict) and isinstance(block.get("text"), str):
                    contents.append(block["text"])

    return contents


def _sanitize_body(body: dict[str, Any], patterns: list[InjectionPattern]) -> dict[str, Any]:
    """Strip matched pattern text from message content in the request body.

    Parameters
    ----------
    body:
        Parsed JSON request body.
    patterns:
        All active patterns to strip.

    Returns
    -------
    dict[str, Any]
        Modified body with matched patterns stripped from message content.
    """
    messages = body.get("messages")
    if not isinstance(messages, list):
        return body

    for msg in messages:
        if not isinstance(msg, dict):
            continue

        content = msg.get("content")
        if isinstance(content, str):
            for pat in patterns:
                content = pat.pattern.sub("", content)
            msg["content"] = content
        elif isinstance(content, list):
            for block in content:
                if isinstance(block, dict) and isinstance(block.get("text"), str):
                    text = block["text"]
                    for pat in patterns:
                        text = pat.pattern.sub("", text)
                    block["text"] = text

    return body


# ---------------------------------------------------------------------------
# Middleware
# ---------------------------------------------------------------------------


class PromptGuardMiddleware(BaseHTTPMiddleware):
    """Starlette middleware for prompt injection detection and defense.

    Intercepts POST requests to LLM proxy endpoints, runs the detection
    pipeline, and applies the configured response mode.

    Parameters
    ----------
    app:
        The ASGI application.
    config:
        PromptGuard configuration. Defaults to warn mode.
    """

    def __init__(self, app: Any, config: PromptGuardConfig | None = None) -> None:
        super().__init__(app)
        self.config = config or PromptGuardConfig()
        self._patterns = list(PATTERNS)

        # Compile custom patterns
        if self.config.custom_patterns:
            self._patterns.extend(compile_custom_patterns(self.config.custom_patterns))

    async def dispatch(self, request: Request, call_next: RequestResponseEndpoint) -> Response:
        """Process the request through the PromptGuard pipeline."""
        # Fast path: disabled or off mode
        if not self.config.enabled or self.config.mode == "off":
            return await call_next(request)

        # Only inspect POST requests to guarded endpoints
        if request.method != "POST" or request.url.path not in _GUARDED_PATHS:
            return await call_next(request)

        # Read and parse the request body
        body_bytes = await request.body()

        try:
            body = json.loads(body_bytes)
        except (json.JSONDecodeError, UnicodeDecodeError):
            logger.warning("PromptGuard: malformed JSON in request body, passing through")
            return await call_next(request)

        if not isinstance(body, dict):
            return await call_next(request)

        # Extract message content
        messages = _extract_messages(body)
        if not messages:
            # No messages to scan -- pass through
            response = await call_next(request)
            response.headers["X-LLMHosts-Threat-Score"] = "0.00"
            response.headers["X-LLMHosts-Threat-Category"] = "none"
            return response

        # Concatenate all messages for scanning
        combined_text = "\n".join(messages)

        # Run detection pipeline
        rule_matches = _detect_rule_based(combined_text, self._patterns)
        heuristic_matches = _detect_heuristic(combined_text)
        all_matches = rule_matches + heuristic_matches

        # Build assessment
        assessment = _build_assessment(all_matches, self.config.mode)

        # Store assessment on request state for downstream handlers
        request.state.threat_assessment = assessment

        # Log assessment
        if self.config.log_all or assessment.threat_score > 0:
            log_level = logging.WARNING if assessment.threat_score >= self.config.threshold else logging.DEBUG
            logger.log(
                log_level,
                "PromptGuard: score=%.2f categories=%s mode=%s details=%s",
                assessment.threat_score,
                ",".join(assessment.categories) or "none",
                assessment.mode,
                assessment.details,
            )

        # Category header value
        category_header = ",".join(assessment.categories) if assessment.categories else "none"

        # Apply mode behavior
        if self.config.mode == "block" and assessment.threat_score >= self.config.threshold:
            return JSONResponse(
                status_code=403,
                content={
                    "error": "prompt_blocked",
                    "threat_score": round(assessment.threat_score, 2),
                    "categories": assessment.categories,
                    "details": assessment.details,
                    "message": "Request blocked by PromptGuard. Configure via [security.prompt_guard] in config.toml",
                },
                headers={
                    "X-LLMHosts-Threat-Score": f"{assessment.threat_score:.2f}",
                    "X-LLMHosts-Threat-Category": category_header,
                },
            )

        if self.config.mode == "sanitize" and assessment.threat_score >= self.config.threshold:
            # Strip matched patterns from the body
            logger.info("PromptGuard: sanitizing request body (score=%.2f)", assessment.threat_score)
            sanitized_body = _sanitize_body(body, self._patterns)
            sanitized_bytes = json.dumps(sanitized_body).encode("utf-8")

            # Reconstruct the request scope with sanitized body
            scope = request.scope

            async def receive() -> dict[str, Any]:
                return {"type": "http.request", "body": sanitized_bytes}

            request = Request(scope, receive)
            request.state.threat_assessment = assessment

        # Forward request and add response headers
        response = await call_next(request)
        response.headers["X-LLMHosts-Threat-Score"] = f"{assessment.threat_score:.2f}"
        response.headers["X-LLMHosts-Threat-Category"] = category_header

        return response
