from __future__ import annotations

import yaml

from leshy.cli.wizard import (
    ALL_CATEGORIES,
    CATEGORY_ATTACKS,
    CATEGORY_LABELS,
    WizardResult,
    generate_custom_test_yaml,
    generate_main_config_yaml,
    generate_test_yaml_for_category,
)


class TestCategoryData:
    def test_all_categories_present(self) -> None:
        expected = {
            "prompt_injection", "jailbreak", "data_extraction",
            "excessive_agency", "safety", "dos",
        }
        assert set(ALL_CATEGORIES) == expected

    def test_all_categories_have_attacks(self) -> None:
        for cat, attacks in CATEGORY_ATTACKS.items():
            assert len(attacks) > 0, f"Category {cat} has no attacks"
            for attack in attacks:
                assert attack.startswith(cat + "."), (
                    f"Attack {attack} doesn't match category {cat}"
                )

    def test_all_categories_have_labels(self) -> None:
        for cat in ALL_CATEGORIES:
            assert cat in CATEGORY_LABELS

    def test_total_attack_count(self) -> None:
        total = sum(len(attacks) for attacks in CATEGORY_ATTACKS.values())
        assert total == 28


class TestGenerateMainConfig:
    def test_basic_config(self) -> None:
        result = WizardResult(
            name="Test Suite",
            llm_model="gpt-4o-mini",
            target_type="http",
            target_url="https://example.com/chat",
            target_body_template='{"message": "{{payload}}"}',
            target_response_path="$.response",
            selected_categories=["prompt_injection", "jailbreak"],
        )
        yaml_str = generate_main_config_yaml(result)

        assert "name: Test Suite" in yaml_str
        assert "llm_model: gpt-4o-mini" in yaml_str
        assert "url: https://example.com/chat" in yaml_str
        assert "tests_dir: leshy/templates" in yaml_str

        # Should be valid YAML
        data = yaml.safe_load(yaml_str)
        assert data["name"] == "Test Suite"
        assert data["tests_dir"] == "leshy/templates"

    def test_with_env_warnings(self) -> None:
        result = WizardResult(
            name="Warn Suite",
            llm_model="claude-sonnet-4-5-20250929",
            target_type="http",
            target_url="https://example.com/chat",
            target_body_template='{"message": "{{payload}}"}',
            selected_categories=["prompt_injection"],
            env_warnings=["ANTHROPIC_API_KEY"],
        )
        yaml_str = generate_main_config_yaml(result)

        assert "# Set before running: export ANTHROPIC_API_KEY=your-key-here" in yaml_str

        data = yaml.safe_load(yaml_str)
        assert data["llm_model"] == "claude-sonnet-4-5-20250929"

    def test_with_headers(self) -> None:
        result = WizardResult(
            name="Auth Suite",
            target_type="http",
            target_url="https://example.com",
            target_headers={"Authorization": "Bearer ${API_KEY}"},
            target_body_template='{"message": "{{payload}}"}',
            selected_categories=["prompt_injection"],
        )
        yaml_str = generate_main_config_yaml(result)
        assert "Authorization" in yaml_str
        assert "${API_KEY}" in yaml_str

    def test_with_preset(self) -> None:
        result = WizardResult(
            name="Preset Suite",
            target_type="http",
            target_url="https://api.openai.com/v1/chat/completions",
            target_preset="openai_chat",
            selected_categories=["prompt_injection"],
        )
        yaml_str = generate_main_config_yaml(result)
        assert "preset: openai_chat" in yaml_str

    def test_langgraph_target(self) -> None:
        result = WizardResult(
            name="LG Suite",
            target_type="langgraph",
            target_url="http://localhost:8123",
            target_headers={"x-assistant-id": "my-agent"},
            selected_categories=["prompt_injection"],
        )
        yaml_str = generate_main_config_yaml(result)
        assert "type: langgraph" in yaml_str
        assert "x-assistant-id" in yaml_str

    def test_custom_loggers(self) -> None:
        result = WizardResult(
            name="Logger Suite",
            target_type="http",
            target_url="https://example.com",
            target_body_template='{"message": "{{payload}}"}',
            loggers=["console", "json_file", "langsmith"],
            selected_categories=["prompt_injection"],
        )
        yaml_str = generate_main_config_yaml(result)
        assert "loggers:" in yaml_str
        assert "- langsmith" in yaml_str
        assert "- json_file" in yaml_str


    def test_with_api_key(self) -> None:
        result = WizardResult(
            name="Key Suite",
            target_type="http",
            target_url="https://example.com",
            target_body_template='{"message": "{{payload}}"}',
            api_key="sk-test-123",
            selected_categories=["prompt_injection"],
        )
        yaml_str = generate_main_config_yaml(result)
        assert "api_key: sk-test-123" in yaml_str
        data = yaml.safe_load(yaml_str)
        assert data["api_key"] == "sk-test-123"

    def test_without_api_key_shows_env_comment(self) -> None:
        result = WizardResult(
            name="Env Suite",
            target_type="http",
            target_url="https://example.com",
            target_body_template='{"message": "{{payload}}"}',
            api_key_env_var="OPENAI_API_KEY",
            selected_categories=["prompt_injection"],
        )
        yaml_str = generate_main_config_yaml(result)
        assert "# api_key: ${OPENAI_API_KEY}" in yaml_str
        # Should be commented out, not parsed as a real key
        data = yaml.safe_load(yaml_str)
        assert "api_key" not in data


class TestGenerateTestYaml:
    def test_generates_valid_yaml(self) -> None:
        for cat in ALL_CATEGORIES:
            yaml_str = generate_test_yaml_for_category(cat)
            data = yaml.safe_load(yaml_str)
            assert "tests" in data
            assert len(data["tests"]) >= 1
            test = data["tests"][0]
            assert "name" in test
            assert "attacks" in test

    def test_includes_all_attacks_for_category(self) -> None:
        yaml_str = generate_test_yaml_for_category("prompt_injection")
        data = yaml.safe_load(yaml_str)
        attack_names = [a["name"] for a in data["tests"][0]["attacks"]]
        expected = CATEGORY_ATTACKS["prompt_injection"]
        assert attack_names == expected

    def test_includes_customization_comments(self) -> None:
        yaml_str = generate_test_yaml_for_category("jailbreak")
        assert "custom_payloads" in yaml_str
        assert "evaluators" in yaml_str
        assert "buffs" in yaml_str


class TestGenerateCustomTestYaml:
    def test_generates_valid_yaml(self) -> None:
        yaml_str = generate_custom_test_yaml()
        data = yaml.safe_load(yaml_str)
        assert "tests" in data
        assert data["tests"] == []

    def test_includes_example_comments(self) -> None:
        yaml_str = generate_custom_test_yaml()
        assert "custom_payloads" in yaml_str
        assert "pattern" in yaml_str
