#!/usr/bin/env python3
"""
Comprehensive test suite for the Alloy language interpreter.
Tests all implemented language features systematically.
"""

import unittest
import sys
import os
from typing import Any, Dict, List, Tuple, Optional
from unittest.mock import Mock, patch

from dotenv import load_dotenv

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

from alloy import Interpreter, parser


class MockInterpreter(Interpreter):
    """Mock interpreter for testing that doesn't make real API calls"""
    
    def __init__(self):
        super().__init__()
        self.openai_client = Mock()
        
    def execute_agent(self, node, env):
        """Return predictable mock responses for agent operations"""
        agent_ref = node[1]
        
        # Mock different responses based on agent type
        if agent_ref == 'alloy':
            mock_response = "mock_alloy_response"
        elif isinstance(agent_ref, tuple) and len(agent_ref) >= 2 and agent_ref[0] == 'agent_var':
            mock_response = f"mock_response_from_{agent_ref[1]}"
        else:
            mock_response = "mock_agent_response"
            
        # Handle different node types, but preserve alloy responses
        if node[0] == 'agent_command':
            command = node[2]
            if isinstance(command, str) and command.startswith('"'):
                command = command.strip('"\'')
            if agent_ref == 'alloy':
                mock_response = f"mock_alloy_response_for_command: {command}"
            else:
                mock_response = f"mock_response_for_command: {command}"
        elif node[0] == 'agent_block':
            statements = node[2]
            if agent_ref == 'alloy':
                mock_response = f"mock_alloy_block_response_with_{len(statements)}_statements"
            else:
                mock_response = f"mock_block_response_with_{len(statements)}_statements"
            
        return mock_response, None


class TestAlloyFeatures(unittest.TestCase):
    """Comprehensive test suite for Alloy language features"""
    
    def setUp(self):
        """Set up test fixtures before each test method"""
        self.interpreter = MockInterpreter()
        self.parser = parser
    
    def parse_and_run(self, code: str) -> Any:
        """Helper method to parse and run code, returning the result"""
        ast = self.parser.parse(code)
        self.assertIsNotNone(ast, "AST should not be None after parsing")
        return self.interpreter.interpret(ast)
    
    def get_var(self, name: str) -> Any:
        """Helper method to get a variable from the global environment"""
        return self.interpreter.global_env.get(name)

    # ========== Core Language Features Tests ==========
    
    def test_variable_assignment_basic(self):
        """Test basic variable assignment with different types"""
        code = '''
        message = "Hello, world!"
        count = 42
        price = 19.99
        is_ready = true
        is_complete = false
        '''
        
        self.parse_and_run(code)
        
        self.assertEqual(self.get_var('message'), "Hello, world!")
        self.assertEqual(self.get_var('count'), 42)
        self.assertEqual(self.get_var('price'), 19.99)
        self.assertEqual(self.get_var('is_ready'), True)
        self.assertEqual(self.get_var('is_complete'), False)
    
    def test_variable_reassignment(self):
        """Test variable reassignment"""
        code = '''
        user_age = 30
        user_age = 31
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('user_age'), 31)
    
    def test_computed_assignment(self):
        """Test assignment with computed values"""
        code = '''
        user_name = "Alice"
        user_age = 30
        full_name = user_name + " Smith"
        is_adult = user_age >= 18
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('full_name'), "Alice Smith")
        self.assertEqual(self.get_var('is_adult'), True)

    # ========== Primitive Types Tests ==========
    
    def test_string_literals(self):
        """Test string literal parsing and assignment"""
        code = '''
        single_quoted = 'Hello world'
        double_quoted = "Hello world"
        empty_string = ""
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('single_quoted'), "Hello world")
        self.assertEqual(self.get_var('double_quoted'), "Hello world")
        self.assertEqual(self.get_var('empty_string'), "")
    
    def test_number_literals(self):
        """Test number literal parsing (integers and floats)"""
        code = '''
        integer = 42
        float_val = 19.99
        negative = -10
        zero = 0
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('integer'), 42)
        self.assertEqual(self.get_var('float_val'), 19.99)
        self.assertEqual(self.get_var('negative'), -10)
        self.assertEqual(self.get_var('zero'), 0)
    
    def test_boolean_literals(self):
        """Test boolean literal parsing"""
        code = '''
        is_true = true
        is_false = false
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('is_true'), True)
        self.assertEqual(self.get_var('is_false'), False)
    
    def test_array_literals(self):
        """Test array literal parsing and assignment"""
        code = '''
        numbers = [1, 2, 3, 4, 5]
        mixed_items = ["hello", 42, true]
        empty_list = []
        nested = [[1, 2], [3, 4]]
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('numbers'), [1, 2, 3, 4, 5])
        self.assertEqual(self.get_var('mixed_items'), ["hello", 42, True])
        self.assertEqual(self.get_var('empty_list'), [])
        self.assertEqual(self.get_var('nested'), [[1, 2], [3, 4]])
    
    def test_object_literals(self):
        """Test object literal parsing and assignment"""
        code = '''
        person = {
            name: "Alice",
            age: 30,
            active: true
        }
        empty_obj = {}
        '''
        
        self.parse_and_run(code)
        expected_person = {"name": "Alice", "age": 30, "active": True}
        self.assertEqual(self.get_var('person'), expected_person)
        self.assertEqual(self.get_var('empty_obj'), {})

    # ========== Arithmetic Operations Tests ==========
    
    def test_basic_arithmetic(self):
        """Test basic arithmetic operations"""
        code = '''
        sum_result = 5 + 3
        difference = 10 - 4
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('sum_result'), 8)
        # Note: The current implementation only supports + operation
        # Additional operations would need to be added to the grammar and interpreter
    
    def test_comparison_operations(self):
        """Test comparison operations"""
        code = '''
        greater = 10 > 5
        greater_equal = 10 >= 10
        less_than = 5 > 10
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('greater'), True)
        self.assertEqual(self.get_var('greater_equal'), True)
        self.assertEqual(self.get_var('less_than'), False)

    # ========== Control Flow Tests ==========
    
    def test_when_then_otherwise_simple(self):
        """Test simple when/then/otherwise statements"""
        code = '''
        user_age = 20
        status = when user_age >= 18 then "adult" otherwise "minor"
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('status'), "adult")
    
    def test_when_then_otherwise_false_condition(self):
        """Test when/then/otherwise with false condition"""
        code = '''
        user_age = 16
        status = when user_age >= 18 then "adult" otherwise "minor"
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('status'), "minor")
    
    def test_when_then_otherwise_block(self):
        """Test when/then/otherwise with block statements"""
        code = '''
        user_age = 25
        user_status = when user_age >= 18 then {
            access_level = "full"
            can_vote = true
            "adult"
        } otherwise {
            access_level = "limited"
            can_vote = false
            "minor"
        }
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('user_status'), "adult")
        # These variables should be local to the block, not global
        # Commenting out until block scope is properly implemented
        # self.assertEqual(self.get_var('access_level'), "full")
        # self.assertEqual(self.get_var('can_vote'), True)
    
    def test_for_loop_simple(self):
        """Test simple for loop with array"""
        code = '''
        items = [1, 2, 3]
        total = 0
        for item in items do {
            total = total + item
        }
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('total'), 6)
    
    def test_while_loop(self):
        """Test while loop"""
        code = '''
        count = 3
        while count > 0 do {
            count = count - 1
        }
        '''
        
        self.parse_and_run(code)
        self.assertEqual(self.get_var('count'), 0)

    # ========== Tool System Tests ==========
    
    def test_tool_definition_basic(self):
        """Test basic tool definition"""
        code = '''
        tool calculate_tax takes amount, rate returns tax_result {
            tax_result = amount * rate
        }
        '''
        
        self.parse_and_run(code)
        # Check that the tool was defined
        self.assertTrue('calculate_tax' in self.interpreter.global_env.tools)
    
    def test_tool_call_basic(self):
        """Test basic tool calling"""
        code = '''
        tool add_numbers takes a, b returns sum {
            sum = a + b
        }
        result = add_numbers(5, 3)
        '''
        
        self.parse_and_run(code)
        # Tools return (result, error) tuples
        result = self.get_var('result')
        self.assertIsInstance(result, tuple)
        self.assertEqual(result[0], 8)  # result
        self.assertIsNone(result[1])    # error
    
    def test_tool_with_multiple_statements(self):
        """Test tool with multiple statements"""
        code = '''
        tool process_order takes order returns order_result {
            tax_amount = order + 0.08
            total = order + tax_amount
            order_result = {
                subtotal: order,
                tax: tax_amount,
                total: total
            }
        }
        result = process_order(100)
        '''
        
        self.parse_and_run(code)
        result = self.get_var('result')
        self.assertIsInstance(result, tuple)
        self.assertIsInstance(result[0], dict)  # Should be an object/dict

    # ========== Tuple Assignment Tests ==========
    
    def test_tuple_assignment_basic(self):
        """Test basic tuple assignment (Result pattern)"""
        code = '''
        tool parse_number takes text returns number_result {
            number_result = (42, null)
        }
        number, error = parse_number("42")
        '''
        
        self.parse_and_run(code)
        # The tool returns (result, error), and tuple assignment should unpack it
        self.assertEqual(self.get_var('number'), 42)
        self.assertIsNone(self.get_var('error'))
    
    def test_tuple_assignment_with_error(self):
        """Test tuple assignment when there's an error"""
        code = '''
        tool failing_tool takes input returns result {
            result = (null, "error occurred")
        }
        value, error = failing_tool("test")
        '''
        
        self.parse_and_run(code)
        self.assertIsNone(self.get_var('value'))
        self.assertEqual(self.get_var('error'), "error occurred")

    # ========== Agent System Tests ==========
    
    def test_clean_agent_definition(self):
        """Test clean agent definition syntax"""
        code = '''
        tool web_search takes query returns search_results {
            search_results = "mock search results for: " + query
        }
        
        tool read_files takes file_path returns file_content {
            file_content = "mock file content from: " + file_path
        }
        
        claude = agent "gpt-4o-mini"
        researcher = agent "gpt-4o-mini" with tools web_search, read_files
        '''
        
        self.parse_and_run(code)
        
        # Check that agents were defined
        claude_info = self.interpreter.global_env.get_agent('claude')
        self.assertEqual(claude_info['model'], 'gpt-4o-mini')
        self.assertEqual(claude_info['tools'], [])
        
        researcher_info = self.interpreter.global_env.get_agent('researcher')
        self.assertEqual(researcher_info['model'], 'gpt-4o-mini')
        self.assertEqual(researcher_info['tools'], ['web_search', 'read_files'])
    
    def test_agent_command_execution(self):
        """Test agent command execution"""
        code = '''
        claude = agent "gpt-4o-mini"
        analysis = claude Analyze "user data"
        '''
        
        self.parse_and_run(code)
        
        # Mock interpreter should return a mock response
        analysis = self.get_var('analysis')
        self.assertIsNotNone(analysis)
        self.assertIn('mock_response', str(analysis))
    
    def test_agent_command_with_tuple_assignment(self):
        """Test agent command with tuple assignment (error handling)"""
        code = '''
        command Analyze means "analyze the provided data"
        claude = agent "gpt-4o-mini"
        analysis, error = claude Analyze "user data"
        '''
        
        self.parse_and_run(code)
        
        # Should have both analysis and error variables
        analysis = self.get_var('analysis')
        error = self.get_var('error')
        self.assertIsNotNone(analysis)
        self.assertIsNone(error)  # Mock returns no error
    
    def test_agent_block_execution(self):
        """Test agent block execution with must/should semantics"""
        code = '''
        command Load means "load and prepare the provided data for analysis"
        command Calculate means "calculate statistical measures and metrics from the data"
        command Provide means "provide insights and trend analysis from calculated data"
        
        analyst = agent "gpt-4o-mini"
        provided_data = "sales_data"
        basic_statistics = "mean, median, mode calculations"
        trend_insights = "quarterly growth patterns"
        
        analysis_result = analyst {
            must Load provided_data
            must Calculate basic_statistics
            should Provide trend_insights
        }
        '''
        
        self.parse_and_run(code)
        
        analysis_result = self.get_var('analysis_result')
        self.assertIsNotNone(analysis_result)
        self.assertIn('mock_block_response', str(analysis_result))
    
    def test_alloy_agent_builtin(self):
        """Test built-in alloy agent"""
        code = '''
        command Review means "review the provided code for improvements"
        code_review = alloy Review "tool_code"
        '''
        
        self.parse_and_run(code)
        
        code_review = self.get_var('code_review')
        self.assertIsNotNone(code_review)
        self.assertIn('mock_alloy_response', str(code_review))

    # ========== Complex Integration Tests ==========
    
    def test_data_processing_pipeline(self):
        """Test a complete data processing pipeline"""
        code = '''command Parse means "parse and structure the provided CSV data for analysis"

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_data(raw_data)
csv_data = cleaned_data
monthly_sales = "January through December sales figures"

analysis = data_analyst { must Parse csv_data }'''
        
        self.parse_and_run(code)
        
        # Check all variables were assigned
        self.assertEqual(self.get_var('raw_data'), "sales_csv_data")
        
        # Check cleaned_data - should be the result from clean_data tool
        cleaned_data = self.get_var('cleaned_data')
        # The tool returns (result, error) tuple, so cleaned_data is a tuple
        self.assertIsInstance(cleaned_data, tuple)
        self.assertEqual(cleaned_data[0], "sales_csv_data_cleaned")
        
        csv_data = self.get_var('csv_data')
        self.assertEqual(csv_data, cleaned_data)
        
        monthly_sales = self.get_var('monthly_sales')
        self.assertEqual(monthly_sales, "January through December sales figures")
        
        analysis = self.get_var('analysis')
        self.assertIsNotNone(analysis)
    
    def test_control_flow_with_agents(self):
        """Test control flow combined with agent operations"""
        code = '''
        command Research means "research the provided topic"
        researcher = agent "gpt-4o-mini"
        
        topic = "machine learning"
        report = when topic == "machine learning" then {
            research_results = researcher Research topic
            research_results
        } otherwise {
            "No research needed"
        }
        '''
        
        self.parse_and_run(code)
        
        report = self.get_var('report')
        self.assertIsNotNone(report)
        # Should contain the mock response from the researcher agent
        self.assertIn('mock_response', str(report))
    
    def test_tools_calling_other_tools(self):
        """Test tools that call other tools"""
        code = '''
        tool helper_tool takes x returns helper_result {
            helper_result = x + 10
        }
        
        tool main_tool takes input returns main_result {
            intermediate, error = helper_tool(input)
            main_result = intermediate + 5
        }
        
        final_result = main_tool(5)
        '''
        
        self.parse_and_run(code)
        
        final_result = self.get_var('final_result')
        self.assertIsInstance(final_result, tuple)
        self.assertEqual(final_result[0], 20)  # 5 + 10 + 5

    # ========== Edge Cases and Boundary Tests ==========
    
    def test_empty_arrays_and_objects(self):
        """Test empty arrays and objects"""
        code = '''
        empty_array = []
        empty_object = {}
        '''
        
        self.parse_and_run(code)
        
        self.assertEqual(self.get_var('empty_array'), [])
        self.assertEqual(self.get_var('empty_object'), {})
    
    def test_nested_data_structures(self):
        """Test nested arrays and objects"""
        code = '''
        nested_data = {
            users: [
                {name: "Alice", age: 30},
                {name: "Bob", age: 25}
            ],
            settings: {
                theme: "dark",
                notifications: true
            }
        }
        '''
        
        self.parse_and_run(code)
        
        nested_data = self.get_var('nested_data')
        self.assertIsInstance(nested_data, dict)
        self.assertIn('users', nested_data)
        self.assertIsInstance(nested_data['users'], list)
        self.assertEqual(len(nested_data['users']), 2)
    
    def test_tool_with_no_parameters(self):
        """Test tool definition with no parameters"""
        code = '''
        tool get_current_time takes returns current_time {
            current_time = "2024-01-01"
        }
        
        time_result = get_current_time()
        '''
        
        self.parse_and_run(code)
        
        time_result = self.get_var('time_result')
        self.assertIsInstance(time_result, tuple)
        self.assertEqual(time_result[0], "2024-01-01")
    
    def test_agent_with_no_tools(self):
        """Test agent definition without tools"""
        code = '''
        simple_agent = agent "gpt-4o-mini"
        response = simple_agent "Hello world"
        '''
        
        self.parse_and_run(code)
        
        response = self.get_var('response')
        self.assertIsNotNone(response)

    # ========== Then Chaining Tests ==========
    
    def test_then_chaining_basic(self):
        """Test basic then chaining with data flow"""
        code = '''
        tool double takes x returns result {
            result = x * 2
        }
        
        tool add_ten takes x returns result {
            result = x + 10
        }
        
        final_result, error = double(5) then add_ten
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify the result - should be (5 * 2) + 10 = 20
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        self.assertEqual(final_result, 20)
        self.assertIsNone(error)
    
    def test_then_chaining_multiple(self):
        """Test multiple then chaining operations"""
        code = '''
        tool double takes x returns result {
            result = x * 2
        }
        
        tool add_ten takes x returns result {
            result = x + 10
        }
        
        tool multiply_by_three takes x returns result {
            result = x * 3
        }
        
        final_result, error = double(5) then add_ten then multiply_by_three
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify the result - should be ((5 * 2) + 10) * 3 = 60
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        self.assertEqual(final_result, 60)
        self.assertIsNone(error)
    
    def test_then_chaining_data_flow(self):
        """Test that data properly flows through chain operations"""
        code = '''
        tool format_name takes name returns result {
            result = "Hello, " + name + "!"
        }
        
        tool make_uppercase takes text returns result {
            result = text  # Mock uppercase for testing
        }
        
        greeting, error = format_name("Alice") then make_uppercase
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify the result flows through
        greeting = self.get_var('greeting')
        error = self.get_var('error')
        
        self.assertEqual(greeting, "Hello, Alice!")
        self.assertIsNone(error)
    
    def test_then_chaining_error_propagation(self):
        """Test error propagation through chain operations"""
        code = '''
        tool failing_tool takes x returns result {
            result = (null, "Tool failed")
        }
        
        tool success_tool takes x returns result {
            result = x + " processed"
        }
        
        final_result, error = failing_tool("test") then success_tool
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify error propagation - should stop at first error
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        self.assertIsNone(final_result)
        self.assertEqual(error, "Tool failed")
    
    def test_then_chaining_with_agents(self):
        """Test chaining operations with agent commands"""
        code = '''
        tool prepare_data takes input returns result {
            result = input + "_prepared"
        }
        
        claude = agent "claude-3.5-sonnet"
        
        final_result, error = prepare_data("raw_data") then claude Analyze
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify mixed chaining works
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        self.assertIsNotNone(final_result)
        self.assertIsNone(error)
        self.assertIn('mock_response', str(final_result))
    
    def test_then_chaining_agent_to_tool(self):
        """Test chaining from agent command to tool"""
        code = '''
        command Research means "research the provided topic"
        tool process_analysis takes analysis returns result {
            result = "Processed: " + analysis
        }
        
        claude = agent "claude-3.5-sonnet"
        
        final_result, error = claude Research "AI trends" then process_analysis
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify agent-to-tool chaining works
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        self.assertIsNotNone(final_result)
        self.assertIsNone(error)
        self.assertIn('Processed:', str(final_result))
    
    def test_then_chaining_complex_pipeline(self):
        """Test complex chaining pipeline with multiple operations"""
        code = '''
        command Process means "process the provided data"

        tool validate_input takes input returns result {
            result = input + "_validated"
        }
        
        tool transform_data takes data returns result {
            result = data + "_transformed"
        }
        
        tool save_result takes processed returns result {
            result = processed + "_saved"
        }
        
        analyst = agent "gpt-4o-mini"
        
        final_result, error = validate_input("user_input") then transform_data then analyst Process then save_result
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify complex pipeline works
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        self.assertIsNotNone(final_result)
        self.assertIsNone(error)
        self.assertIn('_saved', str(final_result))
    
    def test_then_chaining_with_nested_calls(self):
        """Test chaining with nested tool calls"""
        code = '''
        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
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify nested chaining works
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        self.assertIsNotNone(final_result)
        self.assertIsNone(error)
        self.assertIn('_inner_outer_final', str(final_result))
    
    def test_then_chaining_error_in_middle(self):
        """Test error handling when middle operation in chain fails"""
        code = '''
        tool success_tool takes x returns result {
            result = x + "_success"
        }
        
        tool failing_tool takes x returns result {
            result = (null, "Middle operation failed")
        }
        
        final_result, error = success_tool("start") then failing_tool then success_tool
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify error stops chain execution
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        self.assertIsNone(final_result)
        self.assertEqual(error, "Middle operation failed")
    
    def test_then_chaining_single_assignment(self):
        """Test chaining with single variable assignment (no tuple)"""
        code = '''
        tool double takes x returns result {
            result = x * 2
        }
        
        tool add_ten takes x returns result {
            result = x + 10
        }
        
        final_result = double(5) then add_ten
        '''
        
        # Execute the code
        self.parse_and_run(code)
        
        # Verify single assignment works with chaining
        final_result = self.get_var('final_result')
        
        # Should be a tuple (result, error) when not using tuple assignment
        self.assertIsInstance(final_result, tuple)
        self.assertEqual(final_result[0], 20)  # (5 * 2) + 10
        self.assertIsNone(final_result[1])     # No error


class TestAgentBlockExecution(unittest.TestCase):
    """
    Comprehensive test suite for agent block execution features.
    
    These tests follow TDD approach and are designed to FAIL initially
    until agent block functionality is fully implemented.
    
    Features tested:
    - Basic agent block execution
    - Must/Should semantics 
    - System prompt synthesis
    - Implicit variable access
    - Then chaining integration
    - Single command shorthand
    - Error handling
    """
    
    def setUp(self):
        """Set up test fixtures with real API capability"""
        # Use real interpreter for agent tests (will make actual API calls)
        self.interpreter = Interpreter()
        self.parser = parser
        
        # Load API key from .env if available
        load_dotenv()
        
        # Check if we should use mock or real API
        self.use_real_api = os.getenv('OPENAI_API_KEY') is not None
        
        if not self.use_real_api:
            # Fall back to mock for CI/testing without API key
            self.interpreter = MockInterpreter()
    
    def run_code(self, code: str) -> Any:
        """Helper method to parse and run code, returning the result"""
        ast = self.parser.parse(code)
        self.assertIsNotNone(ast, f"Failed to parse code: {code}")
        return self.interpreter.interpret(ast)
    
    def get_var(self, name: str) -> Any:
        """Helper method to get a variable from the global environment"""
        return self.interpreter.global_env.get(name)

    # ========== Basic Agent Block Execution Tests ==========
    
    def test_agent_block_basic_execution(self):
        """Test basic agent block with must commands"""
        code = '''
        command Analyze means "analyze the provided data and return insights"
        
        claude = agent "gpt-4o-mini"
        user_input = "Hello world"
        
        result, error = claude {
            must Analyze user_input
        }
        '''
        
        self.run_code(code)
        
        # Should create variables
        result = self.get_var('result')
        error = self.get_var('error')
        
        # Should make real API call and return response
        self.assertIsNotNone(result)
        self.assertIsInstance(result, str)
        # Error should be None for successful call
        self.assertIsNone(error)
        
        # Result should contain some analysis (basic content check)
        self.assertTrue(len(result) > 0)
    
    def test_agent_block_must_should_semantics(self):
        """Test that must and should commands are handled correctly"""
        code = '''
        command Analyze means "analyze the provided data and return insights"
        command IncludeConfidenceScores means "include confidence scores in the analysis"
        command ProvideReasoning means "provide reasoning for the analysis conclusions"
        
        claude = agent "gpt-4o-mini"
        input_data = "test data for analysis"
        
        result, error = claude {
            must Analyze input_data
            should IncludeConfidenceScores
            should ProvideReasoning
        }
        '''
        
        self.run_code(code)
        
        result = self.get_var('result')
        error = self.get_var('error')
        
        # Should include all commands in single prompt
        self.assertIsNotNone(result)
        self.assertIsNone(error)
        
        # Result should be meaningful response to all commands
        self.assertIsInstance(result, str)
        self.assertTrue(len(result) > 10)  # Should be substantial response
    
    def test_agent_with_system_prompt(self):
        """Test agent definition with system prompt block"""
        code = '''
        command Help_with_python_code means "provide helpful Python coding assistance"
        
        assistant = agent "gpt-4o-mini" with {
            must "Act as a helpful coding assistant"
            should "Provide clear explanations"
        }
        
        result, error = assistant {
            must Help_with_python_code
        }
        '''
        
        self.run_code(code)
        
        result = self.get_var('result')
        error = self.get_var('error')
        
        # Should synthesize system prompt from definition
        self.assertIsNotNone(result)
        self.assertIsNone(error)
        
        # Should reflect the assistant role in response
        self.assertIsInstance(result, str)
        # Response should be coding-oriented based on system prompt
        result_lower = result.lower()
        coding_keywords = ['code', 'python', 'function', 'variable', 'help']
        has_coding_context = any(keyword in result_lower for keyword in coding_keywords)
        self.assertTrue(has_coding_context, "Response should reflect coding assistant role")
    
    def test_agent_block_implicit_variable_access(self):
        """Test that agent blocks can access variables in scope"""
        code = '''
        command Explain means "explain the given topic clearly and comprehensively"
        command AdjustForComplexity means "adjust the explanation for the specified complexity level"
        
        claude = agent "gpt-4o-mini"
        topic = "machine learning"
        complexity = "beginner"
        
        result, error = claude {
            must Explain topic
            should AdjustForComplexity complexity
        }
        '''
        
        self.run_code(code)
        
        result = self.get_var('result')
        error = self.get_var('error')
        
        # Should access topic and complexity variables
        self.assertIsNotNone(result)
        self.assertIsNone(error)
        
        # Should include references to the variables in response
        result_lower = result.lower()
        self.assertIn('machine learning', result_lower)
        self.assertTrue('beginner' in result_lower or 'basic' in result_lower)
    
    def test_agent_block_in_chain(self):
        """Test agent blocks in then chaining"""
        code = '''
        command Analyze means "analyze the provided data and return insights"
        
        tool prepare_text takes text returns result {
            result = text + " (prepared for analysis)"
        }
        
        claude = agent "gpt-4o-mini"
        
        final_result, error = prepare_text("hello") then claude {
            must Analyze
        }
        '''
        
        self.run_code(code)
        
        final_result = self.get_var('final_result')
        error = self.get_var('error')
        
        # Should pass result from prepare_text to claude block
        self.assertIsNotNone(final_result)
        self.assertIsNone(error)
        
        # Should contain analysis of the prepared text
        self.assertIsInstance(final_result, str)
        # Should reference the prepared input
        final_result_lower = final_result.lower()
        prepared_keywords = ['hello', 'prepared', 'analysis']
        has_prepared_context = any(keyword in final_result_lower for keyword in prepared_keywords)
        self.assertTrue(has_prepared_context)
    
    def test_agent_single_command_shorthand(self):
        """Test single command without block syntax"""
        code = '''
        command Analyze means "analyze the provided data and return insights"
        
        claude = agent "gpt-4o-mini"
        text = "hello world"
        
        result, error = claude Analyze text
        '''
        
        self.run_code(code)
        
        result = self.get_var('result')
        error = self.get_var('error')
        
        # Should work without block syntax for single commands
        self.assertIsNotNone(result)
        self.assertIsNone(error)
        
        # Should contain analysis of "hello world"
        self.assertIsInstance(result, str)
        # result_lower = result.lower()
        # self.assertIn('hello', result_lower)
    
    def test_agent_block_error_handling(self):
        """Test proper error handling in agent blocks"""
        code = '''
        command Do_something means "perform some action or task"
        
        invalid_agent = agent "invalid-model-name"
        
        result, error = invalid_agent {
            must Do_something
        }
        '''
        
        self.run_code(code)
        
        result = self.get_var('result')
        error = self.get_var('error')
        
        # Should return proper error tuple
        self.assertIsNone(result)
        self.assertIsNotNone(error)
        self.assertIsInstance(error, str)
        # Error should mention invalid model
        self.assertIn('invalid', error.lower())

    # ========== Advanced Agent Block Features ==========
    
    def test_agent_block_multiple_variables(self):
        """Test agent block accessing multiple variables from scope"""
        code = '''
        command CreateUserProfile means "create a comprehensive user profile"
        command IncludeUserDetails means "include specific user details in the profile"
        
        claude = agent "gpt-4o-mini"
        user_name = "Alice"
        user_age = 25
        user_occupation = "Engineer"
        
        result, error = claude {
            must CreateUserProfile
            should IncludeUserDetails user_name
        }
        '''
        
        self.run_code(code)
        
        result = self.get_var('result')
        error = self.get_var('error')
        
        self.assertIsNotNone(result)
        self.assertIsNone(error)
        
        # Should include all user data in response
        result_lower = result.lower()
        self.assertIn('alice', result_lower)
        self.assertTrue('25' in result or 'twenty' in result_lower)
        self.assertIn('engineer', result_lower)
    
    def test_agent_block_with_complex_data(self):
        """Test agent block with complex data structures"""
        code = '''
        command Analyze means "analyze the provided data and return insights"
        command ProvideRecommendations means "provide recommendations based on the analysis"
        
        claude = agent "gpt-4o-mini"
        user_data = {
            name: "Bob",
            preferences: ["coding", "reading"],
            settings: {
                theme: "dark",
                notifications: true
            }
        }
        
        result, error = claude {
            must Analyze user_data
            should ProvideRecommendations
        }
        '''
        
        self.run_code(code)
        
        result = self.get_var('result')
        error = self.get_var('error')
        
        self.assertIsNotNone(result)
        self.assertIsNone(error)
        
        # Should handle complex data structure
        result_lower = result.lower()
        self.assertIn('bob', result_lower)
        self.assertTrue('coding' in result_lower or 'reading' in result_lower)
    
    def test_agent_block_in_conditional(self):
        """Test agent block within conditional statements"""
        code = '''
        command ProvideAdvancedAnalysis means "provide advanced and detailed analysis"
        command IncludeTechnicalDetails means "include technical details in the analysis"
        
        claude = agent "gpt-4o-mini"
        user_level = "advanced"
        
        analysis = when user_level == "advanced" then {
            result, error = claude {
                must ProvideAdvancedAnalysis
                should IncludeTechnicalDetails
            }
            result
        } otherwise {
            "Basic analysis not needed"
        }
        '''
        
        self.run_code(code)
        
        analysis = self.get_var('analysis')
        
        self.assertIsNotNone(analysis)
        self.assertIsInstance(analysis, str)
        # Should be advanced analysis, not basic message
        self.assertNotEqual(analysis, "Basic analysis not needed")
        
        # Should contain advanced/technical content
        analysis_lower = analysis.lower()
        advanced_keywords = ['advanced', 'technical', 'detailed', 'complex']
        has_advanced_content = any(keyword in analysis_lower for keyword in advanced_keywords)
        self.assertTrue(has_advanced_content)
    
    def test_agent_block_loop_integration(self):
        """Test agent block within loop structures"""
        code = '''
        command Summarize means "provide a concise summary of the given topic"
        command KeepBrief means "keep the response brief and to the point"
        
        claude = agent "gpt-4o-mini"
        topics = ["AI", "ML", "Data"]
        summaries = []
        
        for topic in topics do {
            summary, error = claude {
                must Summarize topic
                should KeepBrief
            }
            summaries = summaries + [summary]
        }
        '''
        
        self.run_code(code)
        
        summaries = self.get_var('summaries')
        
        self.assertIsNotNone(summaries)
        self.assertIsInstance(summaries, list)
        self.assertEqual(len(summaries), 3)
        
        # Each summary should be a string
        for summary in summaries:
            self.assertIsInstance(summary, str)
            self.assertTrue(len(summary) > 0)
    
    def test_agent_block_nested_calls(self):
        """Test nested agent calls within tools"""
        code = '''
        command Analyze means "analyze the provided data and return insights"
        command ProvideInsights means "provide detailed insights from the analysis"
        
        claude = agent "gpt-4o-mini"
        
        tool analyze_and_format takes input_data returns formatted_result {
            analysis, error = claude {
                must Analyze input_data
                should ProvideInsights
            }
            when error then {
                formatted_result = (null, error)
            } otherwise {
                formatted_result = "Report: " + analysis
            }
        }
        
        final_report = analyze_and_format("sales data")
        '''
        
        self.run_code(code)
        
        final_report = self.get_var('final_report')
        
        self.assertIsNotNone(final_report)
        self.assertIsInstance(final_report, tuple)
        
        result, error = final_report
        self.assertIsNotNone(result)
        self.assertIsNone(error)
        self.assertIn("Report:", result)
    
    def test_agent_block_system_prompt_inheritance(self):
        """Test system prompt inheritance in agent definitions"""
        code = '''
        command Help_user means "provide helpful assistance to the user"
        command Analyze_data_trends means "analyze data trends and provide insights"
        
        base_agent = agent "gpt-4o-mini" with {
            must "You are a helpful assistant"
            should "Be concise and accurate"
        }
        
        specialized_agent = agent "gpt-4o-mini" with {
            must "You are a data analyst"
            should "Provide statistical insights"
            should "Use precise terminology"
        }
        
        base_result, base_error = base_agent {
            must Help_user
        }
        
        spec_result, spec_error = specialized_agent {
            must Analyze_data_trends
        }
        '''
        
        self.run_code(code)
        
        base_result = self.get_var('base_result')
        spec_result = self.get_var('spec_result')
        
        self.assertIsNotNone(base_result)
        self.assertIsNotNone(spec_result)
        
        # Results should reflect different system prompts
        base_lower = base_result.lower()
        spec_lower = spec_result.lower()
        
        # Specialized agent should use more analytical language
        analytical_terms = ['data', 'analysis', 'trend', 'statistical', 'metric']
        has_analytical = any(term in spec_lower for term in analytical_terms)
        self.assertTrue(has_analytical)

    # ========== Integration and Complex Scenarios ==========
    
    def test_agent_block_pipeline_processing(self):
        """Test complex pipeline with multiple agent blocks"""
        code = '''
        command Analyze means "analyze the provided data and return insights"
        command IdentifyKeyThemes means "identify the key themes in the data"
        command Summarize means "provide a concise summary of the given content"
        command HighlightMainPoints means "highlight the main points in the summary"
        
        tool validate_input takes input returns result {
            result = input + "_validated"
        }
        
        analyzer = agent "gpt-4o-mini"
        summarizer = agent "gpt-4o-mini" with {
            must "Provide concise summaries"
        }
        
        raw_data = "user feedback data"
        
        # Pipeline: validate -> analyze -> summarize
        validated, val_error = validate_input(raw_data)
        analysis, ana_error = analyzer {
            must Analyze validated
            should IdentifyKeyThemes
        }
        summary, sum_error = summarizer {
            must Summarize analysis
            should HighlightMainPoints
        }
        '''
        
        self.run_code(code)
        
        validated = self.get_var('validated')
        analysis = self.get_var('analysis')
        summary = self.get_var('summary')
        
        self.assertEqual(validated, "user feedback data_validated")
        self.assertIsNotNone(analysis)
        self.assertIsNotNone(summary)
        
        # Summary should be shorter than analysis
        self.assertLess(len(summary), len(analysis))
    
    def test_agent_block_error_recovery(self):
        """Test error recovery in agent block chains"""
        code = '''
        command Process means "process the provided data or request"
        command NoteBackupProcessing means "note that this is backup processing"
        
        primary_agent = agent "gpt-4o-mini"
        backup_agent = agent "gpt-4o-mini"
        
        input_data = "important analysis request"
        
        result, error = primary_agent {
            must Process input_data
        }
        
        final_result = when error then {
            backup_result, backup_error = backup_agent {
                must Process input_data
                should NoteBackupProcessing
            }
            backup_result
        } otherwise {
            result
        }
        '''
        
        self.run_code(code)
        
        final_result = self.get_var('final_result')
        
        # Should have a result either from primary or backup
        self.assertIsNotNone(final_result)
        self.assertIsInstance(final_result, str)

    # ========== Performance and Edge Cases ==========
    
    def test_agent_block_empty_commands(self):
        """Test agent block with empty or minimal commands"""
        code = '''
        command Process means "perform basic processing or analysis"
        
        claude = agent "gpt-4o-mini"
        
        result, error = claude {
            must Process
        }
        '''
        
        self.run_code(code)
        
        result = self.get_var('result')
        error = self.get_var('error')
        
        # Should handle minimal command gracefully
        self.assertIsNotNone(result)
        self.assertIsNone(error)
    
    def test_agent_block_variable_scope_isolation(self):
        """Test variable scope isolation in agent blocks"""
        code = '''
        command AnalyzeVariables means "analyze the provided variables and their scope"
        
        claude = agent "gpt-4o-mini"
        global_var = "global_value"
        
        tool test_tool takes input returns result {
            local_var = "local_value"
            analysis, error = claude {
                must AnalyzeVariables global_var
            }
            result = analysis
        }
        
        tool_result = test_tool("test")
        '''
        
        self.run_code(code)
        
        tool_result = self.get_var('tool_result')
        
        self.assertIsNotNone(tool_result)
        self.assertIsInstance(tool_result, tuple)
        
        result, error = tool_result
        self.assertIsNotNone(result)
        
        # Should access both global and local variables
        result_lower = result.lower()
        self.assertTrue('global' in result_lower or 'global_value' in result_lower)
        self.assertTrue('local' in result_lower or 'local_value' in result_lower)


if __name__ == '__main__':
    # Run the tests
    unittest.main(verbosity=2)