from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Dict, List, Optional

from app.core.config import get_settings


@dataclass
class PreclassificationDecision:
    action: str  # "allow" or "escalate"
    reason: str
    flagged_categories: List[str] = field(default_factory=list)
    scores: Dict[str, float] = field(default_factory=dict)

    @property
    def should_run_stage_two(self) -> bool:
        return self.action != "allow"


_KEYWORDS = {
    "harmfulness": [
        "build a bomb",
        "weapon",
        "kill",
        "harm",
        "explosive",
        "poison",
    ],
    "jailbreak": [
        "ignore previous instructions",
        "jailbreak",
        "break out of your rules",
        "system override",
        "bypass safety",
    ],
    "information_disclosure": [
        "export all company secrets",
        "leak",
        "password list",
        "customer ssn",
    ],
    "implausible_output": [
        "always respond with",
        "replace facts",
    ],
}


def _normalize(text: str) -> str:
    return re.sub(r"\s+", " ", text.lower()).strip()


def run_preclassification(
    *,
    system_prompt: str,
    user_input: str,
    response: str,
    categories: List[str],
) -> PreclassificationDecision:
    settings = get_settings()
    if not getattr(settings, "firewall_preclassifier_enabled", False):
        return PreclassificationDecision(action="escalate", reason="Preclassifier disabled")

    joined = " ".join(filter(None, [system_prompt, user_input, response]))
    normalized = _normalize(joined)
    if not normalized:
        return PreclassificationDecision(action="allow", reason="No content to evaluate")

    flagged: List[str] = []
    scores: Dict[str, float] = {}

    for category in categories:
        patterns = _KEYWORDS.get(category, [])
        score = 0.0
        for keyword in patterns:
            if keyword in normalized:
                score = max(score, 1.0)
                flagged.append(category)
        if score:
            scores[category] = score

    if flagged:
        return PreclassificationDecision(
            action="escalate",
            reason="Potential policy violation detected during preclassification",
            flagged_categories=sorted(set(flagged)),
            scores=scores,
        )

    # If content is short and no risky keywords, allow without stage two
    if len(normalized) < 120:
        return PreclassificationDecision(
            action="allow",
            reason="Content length and heuristics indicate low risk",
            flagged_categories=[],
            scores={category: 0.0 for category in categories},
        )

    return PreclassificationDecision(
        action="escalate",
        reason="Content requires detailed review",
        flagged_categories=[],
        scores={category: 0.0 for category in categories},
    )
