#!/usr/bin/env python3

import sys
import os

# Add the src directory to the Python path
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src'))

from alloy import parser

def test_data_parts():
    
    tests = [
        ("Command definitions", '''
command Parse means "parse and structure the provided CSV data for analysis"
command Calculate means "calculate monthly sales figures and totals from parsed data"
'''),
        
        ("Tool definition", '''
tool clean_data takes raw_data returns cleaned_result {
    cleaned_result = raw_data + "_cleaned"
}
'''),
        
        ("Agent definition", '''
data_analyst = agent "gpt-4o-mini"
'''),
        
        ("Simple assignments", '''
raw_data = "sales_csv_data"
csv_data = cleaned_data
monthly_sales = "January through December sales figures"
'''),
        
        ("Tuple assignment with tool call", '''
cleaned_data, clean_error = clean_data(raw_data)
'''),
        
        ("Agent block", '''
analysis, analysis_error = data_analyst {
    must Parse csv_data
    must Calculate monthly_sales
}
'''),
        
        ("Combined commands + tool", '''
command Parse means "parse"
tool clean_data takes raw_data returns cleaned_result {
    cleaned_result = raw_data + "_cleaned"
}
'''),
        
        ("Combined tool + assignments", '''
tool clean_data takes raw_data returns cleaned_result {
    cleaned_result = raw_data + "_cleaned"
}
raw_data = "sales_csv_data"
'''),
    ]
    
    for name, code in tests:
        print(f"=== {name} ===")
        print(repr(code))
        
        try:
            ast = parser.parse(code)
            if ast and ast[0] == 'program' and len(ast) > 1:
                statements = ast[1]
                print(f"SUCCESS: Parsed {len(statements)} statements")
                for i, stmt in enumerate(statements):
                    if stmt:
                        print(f"  Statement {i}: {stmt[0]}")
                    else:
                        print(f"  Statement {i}: None")
            else:
                print("ERROR: Failed to parse")
                print("AST:", ast)
        except Exception as e:
            print(f"PARSER ERROR: {e}")
        
        print()

if __name__ == "__main__":
    test_data_parts()