#!/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_tool_bodies():
    # Test 1: Two simple assignments
    test1 = '''
tool test1 takes x returns result {
    a = x + 1
    result = a + 2
}
'''
    
    # Test 2: Simple assignment + object assignment
    test2 = '''
tool test2 takes x returns result {
    a = x + 1
    result = {
        value: a
    }
}
'''
    
    # Test 3: Three assignments with object at end
    test3 = '''
tool test3 takes x returns result {
    a = x + 1
    b = a + 2
    result = {
        first: a,
        second: b
    }
}
'''
    
    tests = [
        ("Two simple assignments", test1),
        ("Simple + object assignment", test2), 
        ("Three assignments with object", test3)
    ]
    
    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! Found {len(statements)} statements")
                for i, stmt in enumerate(statements):
                    print(f"  Statement {i}: {stmt[0]}")
                    if stmt[0] == 'tool':
                        tool_body = stmt[4]  # tool body is 5th element
                        if tool_body[0] == 'block':
                            tool_statements = tool_body[1]
                            print(f"    Tool body has {len(tool_statements)} statements:")
                            for j, tool_stmt in enumerate(tool_statements):
                                print(f"      {j}: {tool_stmt[0]} - {tool_stmt}")
            else:
                print("Failed to parse or no statements found")
                print("AST:", ast)
        except Exception as e:
            print(f"Parser error: {e}")
        
        print()

if __name__ == "__main__":
    test_tool_bodies()