"""
Exporting Red Team Test Results

This example demonstrates how to:
1. Export red team test results as a dataset using `red_team_test.export_results_as_dataset`
2. Fetch the exported dataset using `client.datasets.get(dataset_id)`
3. Print red team metrics and export the results as a CSV file

Usage:
    python export_red_team_results.py <test_id>
    python export_red_team_results.py <test_id> -o output.csv
    python export_red_team_results.py <test_id> --output /path/to/output.csv

The CSV file will contain all attacks, judge scores, and other test data from the red team test.
If no output path is specified, it defaults to red_team_results_<test_id>.csv
"""

import argparse
import asyncio
import csv
import os

from haizelabs import AsyncHaize


async def main():
    parser = argparse.ArgumentParser(
        description="Export red team test results as a dataset"
    )
    parser.add_argument(
        "test_id",
        nargs="?",
        help="Red team test ID (optional if RED_TEAM_TEST_ID env var is set)",
    )
    parser.add_argument(
        "-o",
        "--output",
        help="Output CSV file path (default: red_team_results_<test_id>.csv)",
    )
    args = parser.parse_args()

    api_key = os.environ.get("HAIZE_API_KEY")
    if not api_key:
        raise ValueError("Please set HAIZE_API_KEY environment variable")

    test_id = args.test_id or os.environ.get("RED_TEAM_TEST_ID")
    if not test_id:
        raise ValueError(
            "Please provide test_id as an argument or set RED_TEAM_TEST_ID environment variable"
        )

    async with AsyncHaize(api_key=api_key) as haize:

        print("=" * 60)
        print("Red Team Metrics")
        print("=" * 60)
        metrics = await haize.red_team_tests.metrics(test_id)
        print(f"Total Attacks: {metrics.total_attacks}")
        print(f"Average Behavior Duration: {metrics.average_behavior_duration:.2f}s")
        print(
            f"Average Response Length: {metrics.average_response_length:.2f} characters"
        )
        print(f"Exploration Metric: {metrics.exploration_metric:.4f}")
        print(f"Exploration Series Points: {len(metrics.exploration_series)}")
        print(f"Successful Series Points: {len(metrics.num_successful_series)}")
        print("\nFull metrics JSON:")
        print(metrics.model_dump_json(indent=2))

        print(f"\n{'=' * 60}")
        print("Exporting Results as Dataset")
        print("=" * 60)
        dataset_response = await haize.red_team_tests.export_results_as_dataset(
            test_id=test_id,
            name=f"red_team_results_{test_id}",
            description="Exported results from red team test",
            minimum_score=0.0,
        )
        print(f"✓ Exported dataset ID: {dataset_response.dataset_id}")
        print(f"✓ Dataset version: {dataset_response.dataset_version}")
        print(
            f"✓ View in UI: https://platform.haizelabs.com/app/datasets/{dataset_response.dataset_id}"
        )

        print(f"\n{'=' * 60}")
        print("Fetching Exported Dataset")
        print("=" * 60)
        dataset = await haize.datasets.get(dataset_response.dataset_id)
        print(f"✓ Dataset Name: {dataset.dataset_info.name}")
        print(f"✓ Description: {dataset.dataset_info.description}")
        print(f"✓ Number of Rows: {len(dataset.dataset_rows)}")
        print(f"✓ Columns: {[col.name for col in dataset.dataset_info.columns]}")

        output_path = args.output or f"red_team_results_{test_id}.csv"

        if dataset.dataset_rows:
            column_names = [col.formatted_name for col in dataset.dataset_info.columns]
            with open(output_path, "w", newline="", encoding="utf-8") as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow(column_names)
                for row in dataset.dataset_rows:
                    writer.writerow(row.data)

        print(f"\nCSV outputted as {output_path}")


if __name__ == "__main__":
    asyncio.run(main())
