from __future__ import annotations

import json
import logging
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from leshy.core.models import TestSuiteResult, Verdict
from leshy.core.registry import attack_registry

logger = logging.getLogger(__name__)

# Maps attack category prefix → guardrail placement.
CATEGORY_GUARDRAIL_TYPE: dict[str, str] = {
    "prompt_injection": "input",
    "jailbreak": "input",
    "dos": "input",
    "data_extraction": "output",
    "excessive_agency": "output",
    "safety": "output",
}

# Keyword-based fallback patterns for categories that use llm_judge
# (no regex patterns in their default_evaluators).
FALLBACK_PATTERNS: dict[str, list[str]] = {
    "jailbreak": [
        r"(?i)(DAN|do anything now|no restrictions|unrestricted|unfiltered)",
        r"(?i)(jailbreak|bypass|ignore previous|ignore all|override)",
        r"(?i)(pretend you are|act as|roleplay as|you are now)",
    ],
    "safety": [
        r"(?i)(how to make|how to build|instructions for).*(bomb|weapon|drug|poison)",
        r"(?i)(hack into|break into|steal from|rob)",
        r"(?i)(kill|murder|harm|hurt|assault)\s+(a |the |my )",
    ],
}


def _get_category(attack_name: str) -> str:
    """Extract category from dotted attack name, e.g. 'prompt_injection.direct' → 'prompt_injection'."""
    return attack_name.rsplit(".", 1)[0]


def _extract_attack_name_from_payload_id(payload_id: str) -> str | None:
    """Parse attack name from payload_id format 'attack_name:turn:hash'."""
    parts = payload_id.split(":")
    if len(parts) >= 2:
        return parts[0]
    return None


class GuardrailGenerator:
    """Generates LangGraph-compatible guardrail node code from Leshy test results."""

    def __init__(self, results_path: str | Path) -> None:
        self.results_path = Path(results_path)
        self.result = self._load_results()

    @classmethod
    def from_result(cls, result: TestSuiteResult) -> GuardrailGenerator:
        """Create a generator directly from a TestSuiteResult (no file needed)."""
        instance = object.__new__(cls)
        instance.results_path = None
        instance.result = result
        return instance

    def _load_results(self) -> TestSuiteResult:
        data = json.loads(self.results_path.read_text())
        return TestSuiteResult(**data)

    def _find_vulnerable_categories(self) -> dict[str, list[str]]:
        """Return {category: [attack_name, ...]} for test cases with vulnerabilities."""
        categories: dict[str, set[str]] = {}

        for case in self.result.case_results:
            if case.vulnerable_count == 0:
                continue

            # Prefer the attack_names field (populated by runner ≥ v0.2)
            if case.attack_names:
                for attack_name in case.attack_names:
                    cat = _get_category(attack_name)
                    categories.setdefault(cat, set()).add(attack_name)
            else:
                # Fallback: parse from payload_id of vulnerable eval results
                for er in case.eval_results:
                    if er.verdict == Verdict.VULNERABLE:
                        attack_name = _extract_attack_name_from_payload_id(er.payload_id)
                        if attack_name:
                            cat = _get_category(attack_name)
                            categories.setdefault(cat, set()).add(attack_name)

        return {k: sorted(v) for k, v in categories.items()}

    def _get_patterns_for_attack(self, attack_name: str) -> list[str]:
        """Look up regex patterns from the attack's default evaluators."""
        try:
            attack_cls = attack_registry.get(attack_name)
        except KeyError:
            category = _get_category(attack_name)
            return list(FALLBACK_PATTERNS.get(category, []))

        patterns: list[str] = []
        for eval_config in attack_cls.default_evaluators:
            if eval_config.type == "pattern":
                patterns.extend(eval_config.params.get("patterns", []))
            elif eval_config.type == "llm_judge":
                category = _get_category(attack_name)
                patterns.extend(FALLBACK_PATTERNS.get(category, []))

        return patterns

    def generate(self) -> str:
        """Generate the complete guardrail Python file."""
        vuln_categories = self._find_vulnerable_categories()
        if not vuln_categories:
            return self._render_no_vulnerabilities()

        input_patterns: dict[str, list[str]] = {}
        output_patterns: dict[str, list[str]] = {}

        for category, attack_names in vuln_categories.items():
            guardrail_type = CATEGORY_GUARDRAIL_TYPE.get(category, "output")
            patterns: list[str] = []
            for attack_name in attack_names:
                patterns.extend(self._get_patterns_for_attack(attack_name))
            # Deduplicate while preserving order
            patterns = list(dict.fromkeys(patterns))

            if guardrail_type == "input":
                input_patterns[category] = patterns
            else:
                output_patterns[category] = patterns

        return self._render(input_patterns, output_patterns, vuln_categories)

    def write(self, output_path: str | Path) -> Path:
        """Generate and write to file. Returns the output path."""
        path = Path(output_path)
        path.write_text(self.generate())
        return path

    # -- Rendering helpers --------------------------------------------------

    def _render_no_vulnerabilities(self) -> str:
        return (
            '"""\n'
            "LangGraph Guardrail Node — Generated by Leshy\n"
            f"Date: {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M UTC')}\n"
            "\n"
            "No vulnerabilities were detected. No guardrails needed.\n"
            '"""\n'
        )

    def _render(
        self,
        input_patterns: dict[str, list[str]],
        output_patterns: dict[str, list[str]],
        vuln_categories: dict[str, list[str]],
    ) -> str:
        parts: list[str] = []
        parts.append(self._render_header(vuln_categories, input_patterns, output_patterns))
        parts.append(self._render_imports())
        parts.append(self._render_dataclass())
        if input_patterns:
            parts.append(self._render_pattern_dict(input_patterns, "INPUT_PATTERNS", "input"))
        if output_patterns:
            parts.append(self._render_pattern_dict(output_patterns, "OUTPUT_PATTERNS", "output"))
        parts.append(self._render_check_function())
        if input_patterns:
            parts.append(self._render_input_guardrail(list(input_patterns.keys())))
        if output_patterns:
            parts.append(self._render_output_guardrail(list(output_patterns.keys())))
        return "\n\n".join(parts) + "\n"

    def _render_header(
        self,
        vuln_categories: dict[str, list[str]],
        input_patterns: dict[str, list[str]],
        output_patterns: dict[str, list[str]],
    ) -> str:
        total_vulns = sum(
            case.vulnerable_count
            for case in self.result.case_results
            if case.vulnerable_count > 0
        )
        cats = ", ".join(sorted(vuln_categories.keys()))
        timestamp = datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M UTC")

        # Build usage example based on what's generated
        nodes = []
        edges = []
        if input_patterns:
            nodes.append('    graph.add_node("input_guard", input_guardrail)')
            edges.append('    graph.add_edge(START, "input_guard")')
            edges.append('    graph.add_edge("input_guard", "llm")')
        if output_patterns:
            nodes.append('    graph.add_node("output_guard", output_guardrail)')
        if not input_patterns:
            edges.append('    graph.add_edge(START, "llm")')
        if output_patterns:
            edges.append('    graph.add_edge("llm", "output_guard")')
            edges.append('    graph.add_edge("output_guard", END)')
        else:
            edges.append('    graph.add_edge("llm", END)')

        imports = []
        if input_patterns:
            imports.append("input_guardrail")
        if output_patterns:
            imports.append("output_guardrail")

        usage_lines = [
            f"    from guardrail_node import {', '.join(imports)}",
            "",
            "    graph = StateGraph(MessagesState)",
            *nodes,
            '    graph.add_node("llm", call_llm)',
            *edges,
        ]

        return (
            '"""\n'
            "LangGraph Guardrail Nodes — Generated by Leshy Security Testing Framework\n"
            "\n"
            f"Generated from: {self.results_path.name if self.results_path else 'in-memory results'}\n"
            f"Date: {timestamp}\n"
            f"Vulnerabilities: {total_vulns} across {len(vuln_categories)} categories ({cats})\n"
            "\n"
            "Usage:\n"
            + "\n".join(usage_lines)
            + '\n"""'
        )

    def _render_imports(self) -> str:
        return (
            "from __future__ import annotations\n"
            "\n"
            "import re\n"
            "from dataclasses import dataclass\n"
            "\n"
            "try:\n"
            "    from langchain_core.messages import AIMessage\n"
            "except ImportError:\n"
            "\n"
            "    class AIMessage:  # type: ignore[no-redef]\n"
            '        def __init__(self, content: str) -> None:\n'
            "            self.content = content"
        )

    def _render_dataclass(self) -> str:
        return (
            "@dataclass\n"
            "class GuardrailResult:\n"
            '    """Result of a guardrail check."""\n'
            "\n"
            "    blocked: bool\n"
            "    category: str\n"
            "    matched_pattern: str | None = None\n"
            '    message: str = ""'
        )

    def _render_pattern_dict(
        self,
        patterns: dict[str, list[str]],
        var_name: str,
        guardrail_type: str,
    ) -> str:
        direction = "user messages before LLM" if guardrail_type == "input" else "LLM responses before returning"
        lines = [f"# --- Patterns to check {direction} ---"]
        lines.append(f"{var_name}: dict[str, list[re.Pattern[str]]] = {{")
        for category, pattern_list in patterns.items():
            lines.append(f'    "{category}": [')
            for p in pattern_list:
                # Escape backslashes for repr, but keep as raw string
                escaped = p.replace("\\", "\\\\").replace('"', '\\"')
                lines.append(f'        re.compile(r"{escaped}"),')
            lines.append("    ],")
        lines.append("}")
        return "\n".join(lines)

    def _render_check_function(self) -> str:
        return (
            "def _check_patterns(\n"
            "    text: str, patterns: dict[str, list[re.Pattern[str]]]\n"
            ") -> GuardrailResult | None:\n"
            '    """Check text against compiled pattern groups. Returns first match or None."""\n'
            "    for category, compiled_patterns in patterns.items():\n"
            "        for pattern in compiled_patterns:\n"
            "            match = pattern.search(text)\n"
            "            if match:\n"
            "                return GuardrailResult(\n"
            "                    blocked=True,\n"
            "                    category=category,\n"
            "                    matched_pattern=match.group(),\n"
            '                    message=f"Blocked by {category} guardrail",\n'
            "                )\n"
            "    return None"
        )

    def _render_input_guardrail(self, categories: list[str]) -> str:
        cats = ", ".join(categories)
        return (
            "def input_guardrail(state: dict) -> dict:\n"
            f'    """Check incoming messages for security threats.\n'
            "\n"
            "    Place this node BEFORE your LLM node in the graph.\n"
            f"    Detected categories: {cats}\n"
            '    """\n'
            '    messages = state.get("messages", [])\n'
            "    if not messages:\n"
            "        return state\n"
            "\n"
            "    last_message = messages[-1]\n"
            '    content = last_message.content if hasattr(last_message, "content") else str(last_message)\n'
            "\n"
            "    result = _check_patterns(content, INPUT_PATTERNS)\n"
            "    if result and result.blocked:\n"
            "        return {\n"
            '            "messages": [\n'
            "                AIMessage(\n"
            '                    content="I cannot process this request as it may contain harmful content."\n'
            "                )\n"
            "            ]\n"
            "        }\n"
            "    return state"
        )

    def _render_output_guardrail(self, categories: list[str]) -> str:
        cats = ", ".join(categories)
        return (
            "def output_guardrail(state: dict) -> dict:\n"
            f'    """Check outgoing messages for data leaks and policy violations.\n'
            "\n"
            "    Place this node AFTER your LLM node in the graph.\n"
            f"    Detected categories: {cats}\n"
            '    """\n'
            '    messages = state.get("messages", [])\n'
            "    if not messages:\n"
            "        return state\n"
            "\n"
            "    last_message = messages[-1]\n"
            '    content = last_message.content if hasattr(last_message, "content") else str(last_message)\n'
            "\n"
            "    result = _check_patterns(content, OUTPUT_PATTERNS)\n"
            "    if result and result.blocked:\n"
            "        return {\n"
            '            "messages": [\n'
            "                AIMessage(\n"
            "                    content=\"I apologize, but I cannot provide that information\"\n"
            "                    \" as it may contain sensitive data.\"\n"
            "                )\n"
            "            ]\n"
            "        }\n"
            "    return state"
        )
