#!/usr/bin/env python3
"""
Parser-only test suite for the Alloy language.
Tests syntax parsing without executing the interpreter.
This focuses on grammar validation and AST generation.
"""

import unittest
import sys
import os

# 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 AlloyParser, parser


class TestAlloyParser(unittest.TestCase):
    """Test suite focusing on parsing and AST generation"""
    
    def setUp(self):
        """Set up test fixtures"""
        self.parser = parser
    
    def parse_code(self, code: str):
        """Helper method to parse code and return AST"""
        return self.parser.parse(code)
    
    def assert_parses_successfully(self, code: str, description: str = ""):
        """Assert that code parses without errors"""
        try:
            ast = self.parse_code(code)
            self.assertIsNotNone(ast, f"AST should not be None for: {description}")
            return ast
        except Exception as e:
            self.fail(f"Failed to parse {description}: {code}\nError: {e}")

    # ========== Basic Syntax Parsing Tests ==========
    
    def test_parse_variable_assignments(self):
        """Test parsing of basic variable assignments"""
        test_cases = [
            ('name = "Alice"', "string assignment"),
            ('age = 30', "integer assignment"),
            ('price = 19.99', "float assignment"),
            ('is_active = true', "boolean assignment"),
            ('is_complete = false', "boolean assignment"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                ast = self.assert_parses_successfully(code, description)
                self.assertEqual(ast[0], 'program')
                self.assertEqual(ast[1][0][0], 'assign')
    
    def test_parse_array_literals(self):
        """Test parsing of array literals"""
        test_cases = [
            ('numbers = []', "empty array"),
            ('items = [1, 2, 3]', "integer array"),
            ('mixed = ["hello", 42, true]', "mixed type array"),
            ('nested = [[1, 2], [3, 4]]', "nested array"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)
    
    def test_parse_object_literals(self):
        """Test parsing of object literals"""
        test_cases = [
            ('obj = {}', "empty object"),
            ('person = {name: "Alice", age: 30}', "simple object"),
            ('nested = {user: {name: "Bob"}, active: true}', "nested object"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)

    # ========== Tool Definition Parsing Tests ==========
    
    def test_parse_tool_definitions(self):
        """Test parsing of tool definitions"""
        test_cases = [
            ('''tool simple_tool takes x returns result {
                result = x + 1
            }''', "simple tool"),
            
            ('''tool multi_param_tool takes a, b, c returns output {
                output = a + b + c
            }''', "multiple parameter tool"),
            
            ('''tool no_param_tool takes returns value {
                value = 42
            }''', "no parameter tool"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                ast = self.assert_parses_successfully(code, description)
                # Find the tool definition in the AST
                statements = ast[1]
                tool_def = None
                for stmt in statements:
                    if stmt[0] == 'tool':
                        tool_def = stmt
                        break
                self.assertIsNotNone(tool_def, f"Tool definition not found in AST for {description}")
    
    def test_parse_tool_calls(self):
        """Test parsing of tool calls"""
        test_cases = [
            ('result = add(5, 3)', "simple tool call"),
            ('output = process()', "no argument tool call"),
            ('value = complex_tool(x, y, z)', "multiple argument tool call"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)

    # ========== Control Flow Parsing Tests ==========
    
    def test_parse_when_statements(self):
        """Test parsing of when/then/otherwise statements"""
        test_cases = [
            ('status = when age >= 18 then "adult" otherwise "minor"', "simple when statement"),
            
            ('''result = when condition then {
                value = "true branch"
                value
            } otherwise {
                value = "false branch"
                value
            }''', "when statement with blocks"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)
    
    def test_parse_for_loops(self):
        """Test parsing of for loops"""
        test_cases = [
            ('for item in items do process(item)', "simple for loop"),
            
            ('''for item in collection do {
                processed = transform(item)
                results.push(processed)
            }''', "for loop with block"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)
    
    def test_parse_while_loops(self):
        """Test parsing of while loops"""
        test_cases = [
            ('while count < 10 do count = count + 1', "simple while loop"),
            
            ('''while condition do {
                update_state()
                check_condition()
            }''', "while loop with block"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)

    # ========== Agent Syntax Parsing Tests ==========
    
    def test_parse_clean_agent_definitions(self):
        """Test parsing of clean agent definition syntax"""
        test_cases = [
            ('claude = agent "gpt-4o-mini"', "simple agent definition"),
            ('researcher = agent "gpt-4o-mini" with tools web_search, read_files', "agent with tools"),
            ('multi_tool = agent "gpt-4o-mini" with tools tool1, tool2, tool3', "agent with multiple tools"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                ast = self.assert_parses_successfully(code, description)
                # Find the agent definition
                statements = ast[1]
                agent_def = None
                for stmt in statements:
                    if stmt[0] == 'agent_def':
                        agent_def = stmt
                        break
                self.assertIsNotNone(agent_def, f"Agent definition not found for {description}")
    
    def test_parse_agent_commands(self):
        """Test parsing of agent commands"""
        test_cases = [
            ('result = claude "Analyze this data"', "string command"),
            ('analysis = claude Analyze data', "command with data"),
            ('summary = researcher Summarize', "command without data"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)
    
    def test_parse_agent_blocks(self):
        """Test parsing of agent blocks with must/should semantics"""
        test_cases = [
            ('''result = analyst {
                must Load data
                must Calculate statistics
                should Provide insights
            }''', "basic agent block"),
            
            ('''analysis = researcher {
                must Research topic
                should Include examples
            }''', "agent block with data"),
            
            ('''output = processor {
                must Clean input_data
                must Validate results
                should Generate report
                should Include recommendations
            }''', "complex agent block"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)
    
    def test_parse_alloy_agent(self):
        """Test parsing of built-in alloy agent"""
        test_cases = [
            ('review = alloy Review code', "alloy command"),
            ('optimization = alloy Generate tool', "alloy generation"),
            
            ('''reflection = alloy {
                must Analyze input_data
                should Suggest improvements
            }''', "alloy block"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)

    # ========== Tuple Assignment Parsing Tests ==========
    
    def test_parse_tuple_assignments(self):
        """Test parsing of tuple assignments (Result pattern)"""
        test_cases = [
            ('result, error = parse_data(input)', "two-element tuple"),
            ('value, status, message = complex_operation()', "three-element tuple"),
            ('a, b = simple_call()', "simple tuple assignment"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                ast = self.assert_parses_successfully(code, description)
                # Find the tuple assignment
                statements = ast[1]
                tuple_assign = None
                for stmt in statements:
                    if stmt[0] == 'tuple_assign':
                        tuple_assign = stmt
                        break
                self.assertIsNotNone(tuple_assign, f"Tuple assignment not found for {description}")

    # ========== Expression Parsing Tests ==========
    
    def test_parse_binary_operations(self):
        """Test parsing of binary operations"""
        test_cases = [
            ('sum = a + b', "addition"),
            ('is_greater = x > y', "greater than"),
            ('is_valid = count >= minimum', "greater than or equal"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)
    
    def test_parse_chaining_operations(self):
        """Test parsing of chaining with 'then' operator"""
        test_cases = [
            ('result = process_data() then format_output()', "simple chain"),
            ('final = get_data() then validate() then transform()', "multiple chains"),
        ]
        
        for code, description in test_cases:
            with self.subTest(code=code):
                self.assert_parses_successfully(code, description)

    # ========== Complex Program Parsing Tests ==========
    
    def test_parse_complete_program(self):
        """Test parsing of a complete Alloy program"""
        program = '''
        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
        }
        
        # Tool definitions
        tool clean_data takes raw_data returns cleaned_result {
            cleaned_result = raw_data + "_cleaned"
        }
        
        # Agent definitions
        analyst = agent "gpt-4o-mini"
        researcher = agent "gpt-4o-mini" with tools web_search, read_files
        
        # Main program logic
        raw_input = "sales_data.csv"
        cleaned_data, error = clean_data(raw_input)
        
        analysis_result = when error == null then {
            result = analyst {
                must Load cleaned_data
                must Calculate statistics
                should Provide insights
                data cleaned_data
            }
            result
        } otherwise {
            "Data cleaning failed"
        }
        
        # Final processing
        final_report = researcher {
            must Format analysis_result
            should Add recommendations
            data analysis_result
        }
        '''
        
        self.assert_parses_successfully(program, "complete program")
    
    def test_parse_nested_structures(self):
        """Test parsing of deeply nested structures"""
        nested_program = '''
        complex_data = {
            level1: {
                level2: {
                    level3: [
                        {name: "item1", values: [1, 2, 3]},
                        {name: "item2", values: [4, 5, 6]}
                    ]
                }
            }
        }
        
        tool process_nested takes data returns result {
            for item in data.level1.level2.level3 do {
                processed = when item.name == "item1" then {
                    item.values
                } otherwise {
                    []
                }
            }
            result = processed
        }
        '''
        
        self.assert_parses_successfully(nested_program, "nested structures")

    # ========== AST Structure Validation Tests ==========
    
    def test_ast_structure_program(self):
        """Test that AST has correct structure for programs"""
        code = 'x = 42'
        ast = self.parse_code(code)
        
        # AST should be ('program', [statements])
        self.assertEqual(ast[0], 'program')
        self.assertIsInstance(ast[1], list)
        self.assertEqual(len(ast[1]), 1)  # One statement
    
    def test_ast_structure_assignment(self):
        """Test AST structure for assignments"""
        code = 'variable = "value"'
        ast = self.parse_code(code)
        
        statement = ast[1][0]  # First statement
        self.assertEqual(statement[0], 'assign')
        self.assertEqual(statement[1], 'variable')  # Variable name
        self.assertEqual(statement[2][0], 'string')  # Value type
        self.assertEqual(statement[2][1], 'value')  # Value content
    
    def test_ast_structure_tool_definition(self):
        """Test AST structure for tool definitions"""
        code = '''tool test_tool takes a, b returns result {
            result = a + b
        }'''
        ast = self.parse_code(code)
        
        statement = ast[1][0]  # First statement
        self.assertEqual(statement[0], 'tool')
        self.assertEqual(statement[1], 'test_tool')  # Tool name
        self.assertEqual(statement[2], ['a', 'b'])  # Parameters
        self.assertEqual(statement[3], 'result')  # Return variable
        self.assertEqual(statement[4][0], 'block')  # Body is a block
    
    def test_ast_structure_agent_definition(self):
        """Test AST structure for agent definitions"""
        code = 'agent_var = agent "gpt-4o-mini" with tools tool1, tool2'
        ast = self.parse_code(code)
        
        statement = ast[1][0]  # First statement
        self.assertEqual(statement[0], 'agent_def')
        self.assertEqual(statement[1], 'agent_var')  # Agent variable name
        self.assertEqual(statement[2], '"gpt-4o-mini"')  # Model name
        self.assertEqual(statement[3], ['tool1', 'tool2'])  # Tools list

    # ========== Error Recovery Tests ==========
    
    def test_multiple_statements(self):
        """Test parsing multiple statements in sequence"""
        code = '''
        x = 1
        y = 2
        z = x + y
        '''
        
        ast = self.assert_parses_successfully(code, "multiple statements")
        statements = ast[1]
        self.assertEqual(len(statements), 3)
        
        # Check each statement type
        self.assertEqual(statements[0][0], 'assign')
        self.assertEqual(statements[1][0], 'assign')
        self.assertEqual(statements[2][0], 'assign')
    
    def test_empty_program(self):
        """Test parsing empty program"""
        code = ''
        ast = self.parse_code(code)
        
        self.assertEqual(ast[0], 'program')
        self.assertEqual(ast[1], [])  # Empty statement list


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