#!/usr/bin/env python3
import sys
import os
sys.path.insert(0, os.path.join(os.path.dirname(os.path.abspath(__file__)), 'src'))

from alloy import Interpreter, parser
from alloy.environment import Environment

class DebugInterpreter(Interpreter):
    """Debug interpreter that prints variable assignments"""
    
    def __init__(self):
        super().__init__()
        from unittest.mock import Mock
        self.openai_client = Mock()
        
    def execute_agent(self, node, env):
        """Return predictable mock responses for agent operations"""
        print(f"[DEBUG] Executing agent: {node}")
        return "mock_alloy_response", None

    def _visit_chain(self, node, env):
        """Debug version of chain visiting"""
        left_node, right_node = node[1], node[2]
        print(f"[DEBUG] Chain: {left_node} then {right_node}")
        
        # Execute left side first
        left_result = self.visit(left_node, env)
        print(f"[DEBUG] Left result: {left_result}")
        
        # Handle tuple unpacking - some operations return tuples, others don't
        if isinstance(left_result, tuple) and len(left_result) == 2:
            res1, err1 = left_result
        else:
            res1, err1 = left_result, None
        
        print(f"[DEBUG] Unpacked: res1={res1}, err1={err1}")
        
        if err1:
            return None, err1
        
        # Pass result to second operation
        # Handle different types of right-side operations
        if isinstance(right_node, str):
            # Right side is a tool name (e.g., "add_ten")
            # Call the tool with res1 as argument
            print(f"[DEBUG] Calling tool {right_node} with res1={res1}")
            tool = env.get_tool(right_node)
            right_result = tool(res1)
            print(f"[DEBUG] Tool result: {right_result}")
            if isinstance(right_result, tuple) and len(right_result) == 2:
                res2, err2 = right_result
            else:
                res2, err2 = right_result, None
            return res2, err2
        else:
            # Call parent implementation for other cases
            return super()._visit_chain(node, env)
            
    def _visit_tool_def(self, node, env):
        """Debug version of tool definition"""
        name, params, ret_name, body = node[1], node[2], node[3], node[4]
        print(f"[DEBUG] Defining tool {name} with params {params}, returns {ret_name}")
        
        def tool_func(*args):
            print(f"[DEBUG] Tool {name} called with args: {args}")
            tool_env = Environment(env)
            for param, arg in zip(params, args):
                print(f"[DEBUG] Setting {param} = {arg}")
                tool_env.define(param, arg)
            
            # Execute the body statements directly in tool environment
            if isinstance(body, tuple) and body[0] == 'block':
                result = None
                for stmt in body[1]:
                    if isinstance(stmt, str) and stmt == 'return':
                        continue
                    result = self.visit(stmt, tool_env)
            else:
                self.visit(body, tool_env)
            
            # Get the return value from the return variable
            try:
                result = tool_env.get(ret_name)
                print(f"[DEBUG] Tool {name} returning: {result}")
            except:
                result = None
                print(f"[DEBUG] Tool {name} returning None (variable {ret_name} not found)")
            
            return result, None
        
        env.define_tool(name, tool_func)
        return None

# Test individual parts
test_parts = [
    ('command Parse means "parse"', 'command def'),
    ('tool clean_data takes raw_data returns cleaned_result { cleaned_result = raw_data + "_cleaned" }', 'tool def'),
    ('data_analyst = agent "gpt-4o-mini"', 'agent assignment'),
    ('raw_data = "sales_csv_data"', 'simple assignment'),
    ('cleaned_data, clean_error = clean_data(raw_data)', 'tuple assignment'),
]

for code_part, desc in test_parts:
    print(f"\n[DEBUG] Testing {desc}: {code_part}")
    try:
        ast = parser.parse(code_part)
        print(f"[DEBUG] {desc} parsed OK: {ast}")
    except Exception as e:
        print(f"[DEBUG] {desc} FAILED: {e}")

# Test the problematic tool definition
print("[DEBUG] Testing multiline tool definition...")
test_tool = '''tool clean_data takes raw_data returns cleaned_result {
    cleaned_result = raw_data + "_cleaned"
}'''

try:
    ast = parser.parse(test_tool)
    print(f"[DEBUG] Multiline tool parsed OK: {ast}")
except Exception as e:
    print(f"[DEBUG] Multiline tool FAILED: {e}")

# Test one-line version
print("[DEBUG] Testing one-line tool definition...")
test_tool_oneline = 'tool clean_data takes raw_data returns cleaned_result { cleaned_result = raw_data + "_cleaned" }'

try:
    ast = parser.parse(test_tool_oneline)
    print(f"[DEBUG] One-line tool parsed OK: {ast}")
except Exception as e:
    print(f"[DEBUG] One-line tool FAILED: {e}")

# Test commands + tool
print("[DEBUG] Testing commands + tool...")
commands_and_tool = '''command Parse means "parse"
command Calculate means "calculate"

tool clean_data takes raw_data returns cleaned_result { cleaned_result = raw_data + "_cleaned" }'''

try:
    ast = parser.parse(commands_and_tool)
    print(f"[DEBUG] Commands + tool parsed OK: statements={len(ast[1])}")
except Exception as e:
    print(f"[DEBUG] Commands + tool FAILED: {e}")

# Test progressively building up the code
parts = [
    '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 clean_data takes raw_data returns cleaned_result { cleaned_result = raw_data + "_cleaned" }',
    '',
    'data_analyst = agent "gpt-4o-mini"',
    '',
    'raw_data = "sales_csv_data"',
    'cleaned_data, clean_error = clean_data(raw_data)',
    'csv_data = cleaned_data',
    'monthly_sales = "January through December sales figures"',
    '',
    'analysis, analysis_error = data_analyst {',
    '    must Parse csv_data',
    '    must Calculate monthly_sales',
    '}',
]

print("[DEBUG] Testing progressive builds...")
for i in range(1, len(parts) + 1):
    test_code = '\n'.join(parts[:i])
    if test_code.strip():
        print(f"\n[DEBUG] Testing lines 1-{i}:")
        try:
            ast = parser.parse(test_code)
            print(f"[DEBUG] OK - statements: {len(ast[1])}")
        except Exception as e:
            print(f"[DEBUG] FAILED: {e}")
            break

# Test just the problematic sequence
print("[DEBUG] Testing just the sequence that fails...")
problem_sequence = '''command Parse means "parse"
command Calculate means "calculate"

tool clean_data takes raw_data returns cleaned_result { cleaned_result = raw_data + "_cleaned" }

data_analyst = agent "gpt-4o-mini"

raw_data = "sales_csv_data"
cleaned_data, clean_error = clean_data(raw_data)'''

try:
    ast = parser.parse(problem_sequence)
    print(f"[DEBUG] Problem sequence parsed OK: statements={len(ast[1])}")
    for i, stmt in enumerate(ast[1]):
        print(f"[DEBUG] Statement {i}: {stmt}")
except Exception as e:
    print(f"[DEBUG] Problem sequence FAILED: {e}")

# Test if it's the combination causing the issue
print("[DEBUG] Testing without tool definition...")
without_tool = '''command Parse means "parse"
command Calculate means "calculate"

data_analyst = agent "gpt-4o-mini"

raw_data = "sales_csv_data"
cleaned_data, clean_error = clean_data(raw_data)'''

try:
    ast = parser.parse(without_tool)
    print(f"[DEBUG] Without tool parsed OK: statements={len(ast[1])}")
except Exception as e:
    print(f"[DEBUG] Without tool FAILED: {e}")

# Test the exact minimal failing case
print("[DEBUG] Testing minimal failing case...")
minimal_fail = '''cleaned_data, clean_error = clean_data(raw_data)
csv_data = cleaned_data'''

try:
    ast = parser.parse(minimal_fail)
    print(f"[DEBUG] Minimal fail case parsed OK: {ast}")
except Exception as e:
    print(f"[DEBUG] Minimal fail case FAILED: {e}")

# Test with different ordering
print("[DEBUG] Testing different ordering...")
different_order = '''csv_data = cleaned_data
cleaned_data, clean_error = clean_data(raw_data)'''

try:
    ast = parser.parse(different_order)
    print(f"[DEBUG] Different order parsed OK: {ast}")
except Exception as e:
    print(f"[DEBUG] Different order FAILED: {e}")

# Test if it's a variable name issue
print("[DEBUG] Testing with different variable names...")
different_vars = '''x, y = clean_data(raw_data)
z = x'''

try:
    ast = parser.parse(different_vars)
    print(f"[DEBUG] Different vars parsed OK: {ast}")
except Exception as e:
    print(f"[DEBUG] Different vars FAILED: {e}")

# Test the chaining issue
print("[DEBUG] Testing chaining issue...")
chain_test = '''
tool inner_tool takes x returns result {
    result = x + "_inner"
}

tool outer_tool takes x returns result {
    inner_result, inner_error = inner_tool(x)
    when inner_error then {
        result = (null, inner_error)
    } otherwise {
        result = inner_result + "_outer"
    }
}

tool final_tool takes x returns result {
    result = x + "_final"
}

final_result, error = outer_tool("start") then final_tool
'''

try:
    ast = parser.parse(chain_test)
    print(f"[DEBUG] Chain test parsed OK")
    
    interpreter = DebugInterpreter()
    result = interpreter.interpret(ast)
    print("[DEBUG] Chain execution complete")
    
    # Check the result
    try:
        final_result = interpreter.global_env.get('final_result')
        error = interpreter.global_env.get('error')
        print(f"[DEBUG] final_result: {final_result}")
        print(f"[DEBUG] error: {error}")
    except Exception as e:
        print(f"[DEBUG] Variable access error: {e}")
        
except Exception as e:
    print(f"[DEBUG] Chain test FAILED: {e}")
    import traceback
    traceback.print_exc()