#!/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_when_in_tool():
    
    tests = [
        ("Tuple assignment only", '''
tool test takes x returns result {
    a, b = func(x)
}
'''),
        
        ("Tuple assignment + simple assignment", '''
tool test takes x returns result {
    a, b = func(x)
    result = a
}
'''),
        
        ("Tuple assignment + when expression", '''
tool test takes x returns result {
    a, b = func(x)
    when b then result = "error" otherwise result = a
}
'''),
        
        ("When in block form", '''
tool test takes x returns result {
    a, b = func(x)
    when b then {
        result = "error"
    } otherwise {
        result = a
    }
}
'''),
    ]
    
    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]
                if len(statements) > 0 and statements[0][0] == 'tool':
                    tool_body = statements[0][4]
                    if tool_body[0] == 'block':
                        tool_statements = tool_body[1]
                        print(f"SUCCESS: Tool body has {len(tool_statements)} statements")
                        for i, stmt in enumerate(tool_statements):
                            print(f"  Statement {i}: {stmt[0]}")
                    else:
                        print(f"ERROR: Tool body is not a block")
                else:
                    print(f"ERROR: No tool statement found")
            else:
                print("ERROR: Failed to parse")
                print("AST:", ast)
        except Exception as e:
            print(f"PARSER ERROR: {e}")
            import traceback
            traceback.print_exc()
        
        print()

if __name__ == "__main__":
    test_when_in_tool()