"""Generate all 38 new experiment YAML configs for the expanded evaluation."""

from pathlib import Path
import yaml

CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs"

# All 7 models: (file_suffix, groq_model_id)
# New models verified live on Groq as of 2026-04-24 via /openai/v1/models
ALL_MODELS = [
    ("llama8b",        "llama-3.1-8b-instant"),
    ("llama70b",       "llama-3.3-70b-versatile"),
    ("qwen32b",        "qwen/qwen3-32b"),
    ("llama4scout",    "meta-llama/llama-4-scout-17b-16e-instruct"),
    ("gpt_oss_120b",   "openai/gpt-oss-120b"),
    ("gpt_oss_20b",    "openai/gpt-oss-20b"),
    ("groq_compound",  "groq/compound"),
]

# Only the 3 new model families (for existing attack types)
NEW_MODELS = ALL_MODELS[4:]

# hijacking skips llama8b (already done)
HIJACKING_MODELS = ALL_MODELS[1:]

MITIGATION_OPTIONS = {
    "strict_mode": False,
    "sanitize_untrusted_content": True,
    "block_suspicious_tool_actions": True,
    "guard_tool_outputs": True,
    "allow_urls_in_output": False,
}


def build_config(
    name: str,
    description: str,
    model_name: str,
    execution_mode: str,
    selection: dict,
    sample_size: int,
    mitigated: bool,
) -> dict:
    config = {
        "experiment": {
            "name": name,
            "description": description,
            "seed": 42,
            "sample_size": sample_size,
            "output_dir": "results",
        },
        "dataset": {
            "path": "data/normalized/attack_cases.jsonl",
        },
        "selection": selection,
        "model": {
            "provider": "groq",
            "name": model_name,
            "max_tokens": 256,
            "temperature": 0.0,
            "execution_mode": execution_mode,
        },
        "mitigation": {
            "enabled": mitigated,
            "strategy": "surface_aligned" if mitigated else "none",
        },
        "evaluation": {
            "metrics": ["attack_success"],
        },
    }
    if mitigated:
        config["mitigation"]["options"] = MITIGATION_OPTIONS
    return config


def write_config(filename: str, config: dict) -> None:
    path = CONFIGS_DIR / filename
    with path.open("w", encoding="utf-8") as f:
        yaml.dump(config, f, default_flow_style=False, sort_keys=False, allow_unicode=True)
    print(f"  wrote: configs/{filename}")


def model_tag(suffix: str) -> str:
    return f"_{suffix}" if suffix != "llama8b" else ""


def generate_all() -> None:
    CONFIGS_DIR.mkdir(exist_ok=True)
    count = 0

    # -------------------------------------------------------------------------
    # Group 1 — TensorTrust hijacking (6 remaining models, llama8b already done)
    # -------------------------------------------------------------------------
    print("\n[Group 1] TensorTrust hijacking — 6 remaining models × 2 conditions = 12 configs")
    hijacking_selection = {
        "include_sources": ["TensorTrust"],
        "include_attack_subtypes": ["hijacking"],
    }
    for suffix, model_id in HIJACKING_MODELS:
        tag = model_tag(suffix)
        for mitigated in (False, True):
            mitigation_tag = "_mitigated" if mitigated else ""
            name = f"tensortrust_hijacking_25{tag}{mitigation_tag}"
            description = (
                f"TensorTrust hijacking — instruction-override testing "
                f"({'mitigated' if mitigated else 'baseline'}, {suffix})"
            )
            config = build_config(
                name=name,
                description=description,
                model_name=model_id,
                execution_mode="attack_then_task",
                selection=hijacking_selection,
                sample_size=25,
                mitigated=mitigated,
            )
            write_config(f"{name}.yaml", config)
            count += 1

    # -------------------------------------------------------------------------
    # Group 2 — HackAPrompt new families bundled (all 7 models × 2 conditions)
    # -------------------------------------------------------------------------
    print("\n[Group 2] HackAPrompt new families (adaptive+blacklist+xml_escape) — 7 models × 2 = 14 configs")
    new_families_selection = {
        "include_sources": ["HackAPrompt"],
        "include_attack_families": ["adaptive_attack", "blacklist_evasion", "xml_escape_evasion"],
        "require_expected_target": True,
    }
    for suffix, model_id in ALL_MODELS:
        tag = model_tag(suffix)
        for mitigated in (False, True):
            mitigation_tag = "_mitigated" if mitigated else ""
            name = f"hackaprompt_new_families_75{tag}{mitigation_tag}"
            description = (
                f"HackAPrompt: adaptive_attack + blacklist_evasion + xml_escape_evasion — "
                f"{'mitigated' if mitigated else 'baseline'} ({suffix})"
            )
            config = build_config(
                name=name,
                description=description,
                model_name=model_id,
                execution_mode="attack_then_task",
                selection=new_families_selection,
                sample_size=75,
                mitigated=mitigated,
            )
            write_config(f"{name}.yaml", config)
            count += 1

    # -------------------------------------------------------------------------
    # Group 3 — TensorTrust extraction (3 new models × 2 conditions)
    # -------------------------------------------------------------------------
    print("\n[Group 3] TensorTrust extraction — 3 new models × 2 = 6 configs")
    extraction_selection = {
        "include_sources": ["TensorTrust"],
        "include_attack_subtypes": ["extraction"],
    }
    for suffix, model_id in NEW_MODELS:
        for mitigated in (False, True):
            mitigation_tag = "_mitigated" if mitigated else ""
            name = f"tensortrust_extraction_25_{suffix}{mitigation_tag}"
            description = (
                f"TensorTrust extraction — prompt-leakage evaluation "
                f"({'mitigated' if mitigated else 'baseline'}, {suffix})"
            )
            config = build_config(
                name=name,
                description=description,
                model_name=model_id,
                execution_mode="plain",
                selection=extraction_selection,
                sample_size=25,
                mitigated=mitigated,
            )
            write_config(f"{name}.yaml", config)
            count += 1

    # -------------------------------------------------------------------------
    # Group 4 — HackAPrompt override (3 new models × 2 conditions)
    # -------------------------------------------------------------------------
    print("\n[Group 4] HackAPrompt override — 3 new models × 2 = 6 configs")
    override_selection = {
        "include_sources": ["HackAPrompt"],
        "include_attack_families": [
            "direct_instruction_override",
            "compound_instruction_attack",
            "sandwich_defense_bypass",
        ],
        "require_expected_target": True,
    }
    for suffix, model_id in NEW_MODELS:
        for mitigated in (False, True):
            mitigation_tag = "_mitigated" if mitigated else ""
            name = f"hackaprompt_override_25_{suffix}{mitigation_tag}"
            description = (
                f"HackAPrompt instruction-override pressure test "
                f"({'mitigated' if mitigated else 'baseline'}, {suffix})"
            )
            config = build_config(
                name=name,
                description=description,
                model_name=model_id,
                execution_mode="attack_then_task",
                selection=override_selection,
                sample_size=25,
                mitigated=mitigated,
            )
            write_config(f"{name}.yaml", config)
            count += 1

    print(f"\nDone. Generated {count} config files in configs/")


if __name__ == "__main__":
    generate_all()