#!/usr/bin/env python3
import csv
import yaml
import sys

def csv_to_yaml(csv_file, yaml_file):
    """Convert CSV file to YAML format for promptfoo tests"""
    
    tests = []
    
    with open(csv_file, 'r', encoding='utf-8-sig') as f:
        reader = csv.DictReader(f)
        
        for row in reader:
            if row['question']:
                test_case = {
                    'vars': {
                        'question': row['question'].strip()
                    },
                    'assert':[]
                }
                
                if row['regex'] and row['regex'].strip():
                    test_case['assert'].append({
                        'type': 'regex',
                        'value': row['regex'].strip()
                    })

                # Add llm-rubric if present
                if row['llm-rubric'] and row['llm-rubric'].strip():
                    test_case['assert'].append({
                        'type': 'llm-rubric',
                        'value': row['llm-rubric'].strip()
                    })
                
                tests.append(test_case)
    
    with open(yaml_file, 'w', encoding='utf-8') as f:
        yaml.dump(tests, f, default_flow_style=False, allow_unicode=True, sort_keys=False)
    
    print(f"Converted {len(tests)} test cases from {csv_file} to {yaml_file}")

if __name__ == '__main__':
    csv_to_yaml("conversation/tests.csv", "conversation/tests.yaml")
    csv_to_yaml("query/tests.csv", "query/tests.yaml")