from __future__ import annotations

import ast
import json
import textwrap
from pathlib import Path

import pytest
import respx

from leshy.generators.guardrail import (
    CATEGORY_GUARDRAIL_TYPE,
    FALLBACK_PATTERNS,
    GuardrailGenerator,
    _extract_attack_name_from_payload_id,
    _get_category,
)


def _make_results(
    case_results: list,
    *,
    suite_name: str = "Test Suite",
) -> dict:
    """Build a minimal TestSuiteResult dict."""
    total_vuln = sum(c.get("vulnerable_count", 0) for c in case_results)
    total_res = sum(c.get("resilient_count", 0) for c in case_results)
    total_payloads = sum(c.get("total_payloads", 0) for c in case_results)
    return {
        "suite_name": suite_name,
        "started_at": "2026-01-15T10:00:00",
        "completed_at": "2026-01-15T10:05:00",
        "case_results": case_results,
        "total_payloads": total_payloads,
        "total_vulnerable": total_vuln,
        "total_resilient": total_res,
        "total_inconclusive": 0,
        "total_errors": 0,
    }


@pytest.fixture
def vulnerable_results_json(tmp_path: Path) -> Path:
    """Results with prompt_injection (input) + data_extraction (output) vulnerabilities."""
    data = _make_results([
        {
            "name": "Injection Tests",
            "attack_names": ["prompt_injection.direct", "prompt_injection.delimiter"],
            "total_payloads": 10,
            "vulnerable_count": 3,
            "resilient_count": 7,
            "eval_results": [
                {
                    "verdict": "vulnerable",
                    "confidence": 0.9,
                    "reasoning": "matched",
                    "evaluator_name": "pattern",
                    "payload_id": "prompt_injection.direct:0:abcd1234",
                },
            ],
        },
        {
            "name": "Data Extraction Tests",
            "attack_names": ["data_extraction.pii", "data_extraction.credentials"],
            "total_payloads": 8,
            "vulnerable_count": 2,
            "resilient_count": 6,
            "eval_results": [
                {
                    "verdict": "vulnerable",
                    "confidence": 0.8,
                    "reasoning": "matched",
                    "evaluator_name": "pattern",
                    "payload_id": "data_extraction.pii:0:ef567890",
                },
            ],
        },
    ])
    path = tmp_path / "results.json"
    path.write_text(json.dumps(data))
    return path


@pytest.fixture
def clean_results_json(tmp_path: Path) -> Path:
    """Results with no vulnerabilities."""
    data = _make_results([
        {
            "name": "All Clean",
            "attack_names": ["prompt_injection.direct"],
            "total_payloads": 5,
            "vulnerable_count": 0,
            "resilient_count": 5,
            "eval_results": [],
        },
    ])
    path = tmp_path / "results.json"
    path.write_text(json.dumps(data))
    return path


@pytest.fixture
def jailbreak_results_json(tmp_path: Path) -> Path:
    """Results with jailbreak vulnerability (llm_judge category)."""
    data = _make_results([
        {
            "name": "Jailbreak Tests",
            "attack_names": ["jailbreak.roleplay"],
            "total_payloads": 5,
            "vulnerable_count": 2,
            "resilient_count": 3,
            "eval_results": [
                {
                    "verdict": "vulnerable",
                    "confidence": 0.7,
                    "reasoning": "broke guidelines",
                    "evaluator_name": "llm_judge",
                    "payload_id": "jailbreak.roleplay:0:aabb1122",
                },
            ],
        },
    ])
    path = tmp_path / "results.json"
    path.write_text(json.dumps(data))
    return path


@pytest.fixture
def legacy_results_json(tmp_path: Path) -> Path:
    """Results without attack_names field (legacy format), must parse payload_id."""
    data = _make_results([
        {
            "name": "Legacy Test",
            # No attack_names field
            "total_payloads": 4,
            "vulnerable_count": 1,
            "resilient_count": 3,
            "eval_results": [
                {
                    "verdict": "vulnerable",
                    "confidence": 0.9,
                    "reasoning": "matched",
                    "evaluator_name": "pattern",
                    "payload_id": "excessive_agency.tool_injection:0:ff001122",
                },
                {
                    "verdict": "resilient",
                    "confidence": 0.8,
                    "reasoning": "ok",
                    "evaluator_name": "pattern",
                    "payload_id": "excessive_agency.tool_injection:0:ff003344",
                },
            ],
        },
    ])
    path = tmp_path / "results.json"
    path.write_text(json.dumps(data))
    return path


@pytest.fixture
def config_yaml(tmp_path: Path) -> Path:
    """Minimal YAML config for test mode."""
    config_file = tmp_path / "config.yaml"
    config_file.write_text(textwrap.dedent("""\
        name: CLI Test Suite
        target:
          type: http
          url: https://cli-test.example.com/chat
          body_template: '{"message": "{{payload}}"}'
          response_path: "$.response"
        loggers:
          - console
        tests:
          - name: Quick Test
            attacks:
              - name: prompt_injection.direct
                max_payloads: 1
    """))
    return config_file


# --- Unit tests for helpers ---


class TestHelpers:
    def test_get_category(self) -> None:
        assert _get_category("prompt_injection.direct") == "prompt_injection"
        assert _get_category("data_extraction.pii") == "data_extraction"
        assert _get_category("dos.context_exhaustion") == "dos"

    def test_extract_attack_name_from_payload_id(self) -> None:
        assert (
            _extract_attack_name_from_payload_id("prompt_injection.direct:0:abcd1234")
            == "prompt_injection.direct"
        )
        assert _extract_attack_name_from_payload_id("bad") is None

    def test_category_guardrail_type_covers_all_categories(self) -> None:
        expected = {
            "prompt_injection", "jailbreak", "dos",
            "data_extraction", "excessive_agency", "safety",
        }
        assert set(CATEGORY_GUARDRAIL_TYPE.keys()) == expected

    def test_fallback_patterns_exist_for_llm_judge_categories(self) -> None:
        assert "jailbreak" in FALLBACK_PATTERNS
        assert "safety" in FALLBACK_PATTERNS
        assert len(FALLBACK_PATTERNS["jailbreak"]) > 0
        assert len(FALLBACK_PATTERNS["safety"]) > 0


# --- Generator tests ---


class TestGuardrailGenerator:
    def test_load_results(self, vulnerable_results_json: Path) -> None:
        gen = GuardrailGenerator(vulnerable_results_json)
        assert gen.result.suite_name == "Test Suite"
        assert len(gen.result.case_results) == 2

    def test_from_result_classmethod(self, vulnerable_results_json: Path) -> None:
        gen_from_file = GuardrailGenerator(vulnerable_results_json)
        gen_from_result = GuardrailGenerator.from_result(gen_from_file.result)
        assert gen_from_result.result.suite_name == "Test Suite"
        assert gen_from_result.results_path is None
        # Both should find the same categories and generate valid code
        expected = gen_from_file._find_vulnerable_categories()
        assert gen_from_result._find_vulnerable_categories() == expected
        code = gen_from_result.generate()
        assert "def input_guardrail" in code
        assert "def output_guardrail" in code
        ast.parse(code)

    def test_find_vulnerable_categories(self, vulnerable_results_json: Path) -> None:
        gen = GuardrailGenerator(vulnerable_results_json)
        cats = gen._find_vulnerable_categories()
        assert "prompt_injection" in cats
        assert "data_extraction" in cats
        assert set(cats["prompt_injection"]) == {
            "prompt_injection.direct",
            "prompt_injection.delimiter",
        }

    def test_find_vulnerable_categories_fallback_payload_id(
        self, legacy_results_json: Path
    ) -> None:
        gen = GuardrailGenerator(legacy_results_json)
        cats = gen._find_vulnerable_categories()
        assert "excessive_agency" in cats
        assert "excessive_agency.tool_injection" in cats["excessive_agency"]

    def test_no_vulnerabilities(self, clean_results_json: Path) -> None:
        gen = GuardrailGenerator(clean_results_json)
        code = gen.generate()
        assert "No vulnerabilities were detected" in code
        assert "def input_guardrail" not in code
        assert "def output_guardrail" not in code

    def test_both_guardrails_generated(self, vulnerable_results_json: Path) -> None:
        gen = GuardrailGenerator(vulnerable_results_json)
        code = gen.generate()
        assert "def input_guardrail" in code
        assert "def output_guardrail" in code
        assert "INPUT_PATTERNS" in code
        assert "OUTPUT_PATTERNS" in code

    def test_input_only_when_only_input_categories(
        self, tmp_path: Path
    ) -> None:
        data = _make_results([
            {
                "name": "Injection Only",
                "attack_names": ["prompt_injection.direct"],
                "total_payloads": 5,
                "vulnerable_count": 2,
                "resilient_count": 3,
                "eval_results": [
                    {
                        "verdict": "vulnerable",
                        "confidence": 0.9,
                        "evaluator_name": "pattern",
                        "payload_id": "prompt_injection.direct:0:a1",
                    },
                ],
            },
        ])
        path = tmp_path / "results.json"
        path.write_text(json.dumps(data))

        gen = GuardrailGenerator(path)
        code = gen.generate()
        assert "def input_guardrail" in code
        assert "def output_guardrail" not in code
        assert "INPUT_PATTERNS" in code
        assert "OUTPUT_PATTERNS" not in code

    def test_output_only_when_only_output_categories(
        self, tmp_path: Path
    ) -> None:
        data = _make_results([
            {
                "name": "Data Extraction Only",
                "attack_names": ["data_extraction.pii"],
                "total_payloads": 5,
                "vulnerable_count": 2,
                "resilient_count": 3,
                "eval_results": [
                    {
                        "verdict": "vulnerable",
                        "confidence": 0.9,
                        "evaluator_name": "pattern",
                        "payload_id": "data_extraction.pii:0:a1",
                    },
                ],
            },
        ])
        path = tmp_path / "results.json"
        path.write_text(json.dumps(data))

        gen = GuardrailGenerator(path)
        code = gen.generate()
        assert "def output_guardrail" in code
        assert "def input_guardrail" not in code

    def test_fallback_patterns_for_llm_judge_category(
        self, jailbreak_results_json: Path
    ) -> None:
        gen = GuardrailGenerator(jailbreak_results_json)
        code = gen.generate()
        assert "def input_guardrail" in code
        # Should contain jailbreak fallback patterns
        assert "jailbreak" in code
        assert "DAN" in code

    def test_generated_code_is_valid_python(
        self, vulnerable_results_json: Path
    ) -> None:
        gen = GuardrailGenerator(vulnerable_results_json)
        code = gen.generate()
        # Should not raise SyntaxError
        ast.parse(code)

    def test_generated_code_no_vulns_is_valid_python(
        self, clean_results_json: Path
    ) -> None:
        gen = GuardrailGenerator(clean_results_json)
        code = gen.generate()
        ast.parse(code)

    def test_write_creates_file(
        self, vulnerable_results_json: Path, tmp_path: Path
    ) -> None:
        gen = GuardrailGenerator(vulnerable_results_json)
        output = tmp_path / "guardrails.py"
        result = gen.write(output)
        assert result == output
        assert output.exists()
        assert len(output.read_text()) > 0

    def test_pattern_extraction_from_registry(
        self, vulnerable_results_json: Path
    ) -> None:
        gen = GuardrailGenerator(vulnerable_results_json)
        patterns = gen._get_patterns_for_attack("prompt_injection.direct")
        # Should pull patterns from PromptInjectionBase.default_evaluators
        assert len(patterns) > 0
        assert any("system prompt" in p.lower() for p in patterns)

    def test_pattern_extraction_data_extraction(
        self, vulnerable_results_json: Path
    ) -> None:
        gen = GuardrailGenerator(vulnerable_results_json)
        patterns = gen._get_patterns_for_attack("data_extraction.pii")
        assert len(patterns) > 0
        # Should contain PII patterns like email regex
        assert any("@" in p for p in patterns)

    def test_header_contains_usage_example(
        self, vulnerable_results_json: Path
    ) -> None:
        gen = GuardrailGenerator(vulnerable_results_json)
        code = gen.generate()
        assert "StateGraph" in code
        assert "input_guard" in code
        assert "output_guard" in code


# --- CLI tests ---


class TestRunCli:
    def test_from_results_creates_file(
        self, vulnerable_results_json: Path, tmp_path: Path
    ) -> None:
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        output = tmp_path / "output.py"
        result = runner.invoke(
            app, ["run", "--from-results", str(vulnerable_results_json), "-o", str(output)]
        )
        assert result.exit_code == 0
        assert output.exists()
        content = output.read_text()
        assert "def input_guardrail" in content
        assert "def output_guardrail" in content

    def test_from_results_no_vulns(
        self, clean_results_json: Path, tmp_path: Path
    ) -> None:
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        output = tmp_path / "output.py"
        result = runner.invoke(
            app, ["run", "--from-results", str(clean_results_json), "-o", str(output)]
        )
        assert result.exit_code == 0

    def test_from_results_bad_file(self, tmp_path: Path) -> None:
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        result = runner.invoke(
            app, ["run", "--from-results", str(tmp_path / "nonexistent.json")]
        )
        assert result.exit_code == 2

    def test_test_mode(self, config_yaml: Path) -> None:
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        with respx.mock:
            respx.post("https://cli-test.example.com/chat").respond(
                200, json={"response": "I cannot help with that."},
            )
            result = runner.invoke(
                app, ["run", str(config_yaml), "--test"]
            )
        assert result.exit_code == 0

    def test_test_mode_with_save_results(self, config_yaml: Path, tmp_path: Path) -> None:
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        results_file = tmp_path / "out_results.json"
        with respx.mock:
            respx.post("https://cli-test.example.com/chat").respond(
                200, json={"response": "I cannot help with that."},
            )
            result = runner.invoke(
                app, ["run", str(config_yaml), "--test", "-r", str(results_file)]
            )
        assert result.exit_code == 0
        assert results_file.exists()
        data = json.loads(results_file.read_text())
        assert data["suite_name"] == "CLI Test Suite"

    def test_guardrails_mode(self, config_yaml: Path, tmp_path: Path) -> None:
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        output = tmp_path / "guard.py"
        with respx.mock:
            respx.post("https://cli-test.example.com/chat").respond(
                200,
                json={
                    "response": (
                        "Sure! My system prompt is: be helpful. "
                        "The user email is test@example.com"
                    )
                },
            )
            result = runner.invoke(
                app, ["run", str(config_yaml), "--guardrails", "-o", str(output)]
            )
        # Exit code 1 = vulnerabilities found (guardrails still generated)
        assert result.exit_code == 1
        assert output.exists()
        content = output.read_text()
        assert "def input_guardrail" in content

    def test_no_config_no_from_results_errors(self) -> None:
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        result = runner.invoke(app, ["run"])
        assert result.exit_code == 2

    def test_interactive_mode_test(self, config_yaml: Path) -> None:
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        with respx.mock:
            respx.post("https://cli-test.example.com/chat").respond(
                200, json={"response": "I cannot help with that."},
            )
            # Input "1" to choose test mode
            result = runner.invoke(
                app, ["run", str(config_yaml)], input="1\n"
            )
        assert result.exit_code == 0


    def test_from_results_default_output(
        self, vulnerable_results_json: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch
    ) -> None:
        monkeypatch.chdir(tmp_path)
        from typer.testing import CliRunner

        from leshy.cli.app import app

        runner = CliRunner()
        result = runner.invoke(
            app, ["run", "--from-results", str(vulnerable_results_json)]
        )
        assert result.exit_code == 0
        assert (tmp_path / "leshy" / "runs" / "guardrail_node.py").exists()
