#!/usr/bin/env python3
"""Run compute_metrics over every results folder in the results directory.

This script iterates through all model/dataset/run combinations in the results
directory and computes metrics for each one.
"""

import json
from pathlib import Path
from typing import Any

from compute_metrics import Metrics


JAILBREAK_DATASET_CLEAN_NAME = "centrepourlasecuriteia-jailbreak-dataset"

JAILBREAK_DATASETS = {
    "allenai-wildjailbreak",
    "bells-o-project-jailbreak-dataset",
    "centrepourlasecuriteia-jailbreak-dataset",
    "bells-o-project-BET-jailbreak-dataset",
    "centrepourlasecuriteia-BET-jailbreak-dataset",
    "deepset-prompt-injections",
    "local_datasets-malignant",
    "TrustAIRLab-in-the-wild-jailbreak-prompts",
}

CONTENT_MODERATION_DATASETS = {
    "bells-o-project-content-moderation-input",
    "centrepourlasecuriteia-content-moderation-input-dataset",
    "bells-o-project-content-moderation-output",
    "centrepourlasecuriteia-content-moderation-output-dataset",
}


def build_technique_map() -> dict[str, str]:
    """Build a mapping from prompt_id to technique_type by loading the HF jailbreak dataset.

    Returns:
        dict mapping prompt_id → technique_type

    """
    from hashlib import sha256

    from datasets import load_dataset

    ds = load_dataset("centrepourlasecuriteia/jailbreak-dataset")
    technique_map = {}
    for sample in ds["train"]:
        hash_id = sha256(sample["prompt"].encode()).hexdigest()
        prompt_id = f"{JAILBREAK_DATASET_CLEAN_NAME}_{hash_id}"
        technique_map[prompt_id] = sample["technique_type"]
    return technique_map


def find_all_result_combinations(results_dir: Path, eval_type: str = "jailbreak") -> list[tuple[str, str, str]]:
    """Find all model_provider_use_case/dataset/model_name combinations in the results directory.

    Args:
        results_dir: Base results directory.
        eval_type: Type of evaluation ("jailbreak" or "content_moderation").
            Only dataset directories matching the eval type are included.

    Returns:
        List of (model_provider_use_case, dataset_name, model_name) tuples.

    """
    allowed_datasets = JAILBREAK_DATASETS if eval_type == "jailbreak" else CONTENT_MODERATION_DATASETS
    combinations = []

    if not results_dir.exists():
        return combinations

    # Iterate through all model_provider_use_case directories
    for provider_dir in sorted(results_dir.iterdir()):
        if not provider_dir.is_dir():
            continue

        model_provider_use_case = provider_dir.name

        # Iterate through all dataset directories within the provider directory
        for dataset_dir in sorted(provider_dir.iterdir()):
            if not dataset_dir.is_dir():
                continue

            dataset_name = dataset_dir.name

            if dataset_name not in allowed_datasets:
                continue

            # Iterate through all model directories
            for model_dir in sorted(dataset_dir.iterdir()):
                if not model_dir.is_dir():
                    continue

                # Check if this directory contains JSON files
                json_files = list(model_dir.glob("*.json"))
                if json_files:
                    model_name = model_dir.name
                    combinations.append((model_provider_use_case, dataset_name, model_name))

    return combinations


def compute_all_metrics(
    results_dir: Path,
    output_file: Path | None = None,
    eval_type: str = "jailbreak",
) -> dict[str, Any]:
    """Compute metrics for all result combinations.

    Args:
        results_dir: Base results directory.
        output_file: Optional path to save results as JSON.
        eval_type: Type of evaluation ("jailbreak" or "content_moderation").

    Returns:
        Dictionary mapping (model_provider_use_case, dataset_name, model_name) to metrics.

    """
    mapping_file = Path("data/model_info_mapping.json").resolve()
    metrics_calculator = Metrics(results_dir, mapping_file=mapping_file)
    all_combinations = find_all_result_combinations(results_dir, eval_type=eval_type)

    # Build technique map once if processing jailbreak results
    technique_map = None
    if eval_type == "jailbreak":
        has_jailbreak = any(ds == JAILBREAK_DATASET_CLEAN_NAME for _, ds, _ in all_combinations)
        if has_jailbreak:
            print("Loading jailbreak dataset to build technique_type map...")
            technique_map = build_technique_map()
            print(f"  Built technique map with {len(technique_map)} entries\n")

    results = {}

    print(f"Found {len(all_combinations)} result combinations to process\n")

    for i, (model_provider_use_case, dataset_name, model_name) in enumerate(all_combinations, 1):
        print(f"[{i}/{len(all_combinations)}] Processing: {model_provider_use_case}/{dataset_name}/{model_name}")

        metrics = metrics_calculator.compute_all_metrics(
            model_provider_use_case=model_provider_use_case, dataset_name=dataset_name, model_name=model_name
        )

        # Add per-technique accuracy for jailbreak dataset
        if technique_map is not None and dataset_name == JAILBREAK_DATASET_CLEAN_NAME:
            metrics["accuracy_per_technique"] = metrics_calculator.compute_accuracy_per_technique(technique_map)

        key = f"{model_provider_use_case}/{dataset_name}/{model_name}"
        results[key] = {
            "model_provider_use_case": model_provider_use_case,
            "dataset_name": dataset_name,
            "model_name": model_name,
            **metrics,
        }
        if isinstance(metrics["fpr"], float):
            fpr = f"{metrics['fpr']:.4f}"
        else:
            fpr = "N/A"
        print(f"  ✓ Accuracy: {metrics['accuracy']:.4f}, FPR: {fpr}, Samples: {metrics['num_samples']}")

    if output_file:
        print(f"\nSaving results to {output_file}")
        with open(output_file, "w") as f:
            json.dump(results, f, indent=2)
        print(f"✓ Results saved to {output_file}")

    return results


def main():
    """Main entry point."""
    import argparse

    parser = argparse.ArgumentParser(description="Compute metrics for all result folders in the results directory")
    parser.add_argument(
        "--results-dir", type=Path, default=Path("results"), help="Base results directory (default: results)"
    )
    parser.add_argument("--output", type=Path, help="Optional JSON file to save all results")
    parser.add_argument(
        "--type",
        type=str,
        choices=["jailbreak", "content_moderation"],
        default="jailbreak",
        help="Evaluation type (default: jailbreak). Sets default output path if --output is not specified.",
    )

    args = parser.parse_args()

    # Set default output path based on type if not explicitly provided
    output_file = args.output
    if output_file is None:
        if args.type == "jailbreak":
            output_file = Path("data/jailbreak_metrics.json")
        else:
            output_file = Path("data/content_moderation_metrics.json")

    results = compute_all_metrics(args.results_dir, output_file, eval_type=args.type)

    print(f"\n{'=' * 60}")
    print(f"Summary: Processed {len(results)} result combinations")
    print(f"{'=' * 60}")


if __name__ == "__main__":
    main()