#!/usr/bin/env python
"""
Generate synthetic test cases using Strands ExperimentGenerator.

This script replaces the custom generate_data.py and uses the SDK
to generate diverse, high-quality test cases for the fraud agent.
"""
import os
import json
import asyncio
from dotenv import load_dotenv

load_dotenv()

from typing import List, Dict
from strands_evals.generators import ExperimentGenerator
from strands_evals.evaluators import OutputEvaluator
from evals.config import eval_model

# Context description for the generator
CONTEXT = """
You are generating test cases for a Fraud Model Explainability Assistant for a financial services company.
The assistant uses RAG and tools to explain fraud scores (0-1000), SHAP values, and compliance checks.

Users are typically:
1. Fraud Analysts (investigating specific cases)
2. Data Scientists (monitoring model performance)
3. Compliance Officers (checking for Fair Lending bias)
4. Executives (asking for high-level summaries)

Tools available:
- get_application_summary(app_id): Returns score, risk level.
- explain_fraud_score(app_id): Returns SHAP feature contributions.
- compare_to_population(app_id): Returns stats vs approved/denied.
- check_fair_lending_flags(app_id): Returns bias analysis.
- get_identity_network(app_id): Returns linked applications.
"""

async def generate():
    print("🚀 Starting Experiment Generation with SDK...")
    
    # Initialize generator with str input/output
    generator = ExperimentGenerator[str, str](
        input_type=str,
        output_type=str,
        model=eval_model
    )
    
    # Generate experiment
    print("   Generating cases (this may take a minute)...")
    experiment = await generator.from_context_async(
        context=CONTEXT,
        num_cases=10,  # Generate 10 new cases
        evaluator=OutputEvaluator, # Pass class, let generator create rubric
        task_description="Explain fraud model decisions and risk factors.",
        num_topics=5   # Split across different topics (High Risk, Compliance, etc.)
    )
    
    print(f"✅ Generated {len(experiment.cases)} new test cases.")
    
    # Convert to our JSON format
    new_cases = []
    for i, case in enumerate(experiment.cases):
        # Metadata might be None
        metadata = case.metadata if case.metadata else {}
        new_case = {
            "id": f"synth_sdk_{i+1}",
            "question": case.input,
            "expected_intent": metadata.get("topic", "General"),
            "expected_answer_key_points": [case.expected_output] if case.expected_output else []
        }
        new_cases.append(new_case)
        print(f"   - [{new_case['expected_intent']}] {new_case['question'][:60]}...")

    # Load existing cases to append (optional, or overwrite)
    output_path = "evaluation/dataset_sdk.json"
    
    # Saving to a new file to avoid overwriting the main dataset during this test
    with open(output_path, "w") as f:
        json.dump(new_cases, f, indent=2)
        
    print(f"\n💾 Saved {len(new_cases)} cases to {output_path}")
    print("   Review the file and merge into evaluation/dataset.json if desired.")

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