"""
AST Node Definitions for the Alloy Programming Language

This module defines the abstract syntax tree (AST) node types used by the
Alloy parser and interpreter. Each node type represents a different construct
in the Alloy language.
"""

from typing import Any, List, Dict, Optional, Union, Tuple
from abc import ABC, abstractmethod


class ASTNode(ABC):
    """Base class for all AST nodes"""
    
    def __init__(self, node_type: str):
        self.node_type = node_type


class Program(ASTNode):
    """Root node representing an entire Alloy program"""
    
    def __init__(self, statements: List[ASTNode]):
        super().__init__("program")
        self.statements = statements


class Statement(ASTNode):
    """Base class for all statement nodes"""
    pass


class Expression(ASTNode):
    """Base class for all expression nodes"""
    pass


# Literals
class StringLiteral(Expression):
    """String literal expression"""
    
    def __init__(self, value: str):
        super().__init__("string")
        self.value = value


class NumberLiteral(Expression):
    """Number literal expression"""
    
    def __init__(self, value: Union[int, float]):
        super().__init__("number")
        self.value = value


class BooleanLiteral(Expression):
    """Boolean literal expression"""
    
    def __init__(self, value: bool):
        super().__init__("bool")
        self.value = value


class ArrayLiteral(Expression):
    """Array literal expression"""
    
    def __init__(self, elements: List[Expression]):
        super().__init__("array")
        self.elements = elements


class ObjectLiteral(Expression):
    """Object literal expression"""
    
    def __init__(self, pairs: List[Tuple[str, Expression]]):
        super().__init__("object")
        self.pairs = pairs


# Variables and Assignment
class Identifier(Expression):
    """Variable identifier"""
    
    def __init__(self, name: str):
        super().__init__("identifier")
        self.name = name


class Assignment(Statement):
    """Variable assignment statement"""
    
    def __init__(self, name: str, value: Expression):
        super().__init__("assign")
        self.name = name
        self.value = value


class TupleAssignment(Statement):
    """Tuple assignment statement (e.g., result, error = tool_call())"""
    
    def __init__(self, names: List[str], value: Expression):
        super().__init__("tuple_assign")
        self.names = names
        self.value = value


# Tool Definitions
class ToolDefinition(Statement):
    """Tool (function) definition"""
    
    def __init__(self, name: str, params: List[str], return_name: str, body: 'Block'):
        super().__init__("tool")
        self.name = name
        self.params = params
        self.return_name = return_name
        self.body = body


class ToolCall(Expression):
    """Tool (function) call expression"""
    
    def __init__(self, name: str, args: List[Expression]):
        super().__init__("tool_call")
        self.name = name
        self.args = args


# Agent Definitions and Operations
class AgentDefinition(Statement):
    """Agent definition statement"""
    
    def __init__(self, name: str, model: str, tools: List[str]):
        super().__init__("agent_def")
        self.name = name
        self.model = model
        self.tools = tools


class AgentReference(Expression):
    """Reference to an agent"""
    
    def __init__(self, ref_type: str, value: str = None):
        super().__init__("agent_ref")
        self.ref_type = ref_type  # 'alloy', 'agent_var', or 'agent'
        self.value = value


class AgentCommand(Expression):
    """Agent command execution"""
    
    def __init__(self, agent_ref: AgentReference, command: Union[str, Expression], data: Optional[Expression] = None):
        super().__init__("agent_command")
        self.agent_ref = agent_ref
        self.command = command
        self.data = data


class AgentBlock(Expression):
    """Agent block with must/should statements"""
    
    def __init__(self, agent_ref: AgentReference, statements: List['AgentStatement']):
        super().__init__("agent_block")
        self.agent_ref = agent_ref
        self.statements = statements


class AgentStatement(ASTNode):
    """Statement within an agent block (must, should, data)"""
    
    def __init__(self, statement_type: str, content: Union[str, Expression]):
        super().__init__(statement_type)
        self.content = content


# Control Flow
class Block(ASTNode):
    """Block of statements enclosed in braces"""
    
    def __init__(self, statements: List[Statement]):
        super().__init__("block")
        self.statements = statements


class WhenStatement(Statement):
    """Conditional when/then/otherwise statement"""
    
    def __init__(self, condition: Expression, then_stmt: Union[Statement, Block], otherwise_stmt: Union[Statement, Block]):
        super().__init__("when")
        self.condition = condition
        self.then_stmt = then_stmt
        self.otherwise_stmt = otherwise_stmt


class ForLoop(Statement):
    """For loop statement"""
    
    def __init__(self, variable: str, iterable: Expression, body: Union[Statement, Block]):
        super().__init__("for")
        self.variable = variable
        self.iterable = iterable
        self.body = body


class WhileLoop(Statement):
    """While loop statement"""
    
    def __init__(self, condition: Expression, body: Union[Statement, Block]):
        super().__init__("while")
        self.condition = condition
        self.body = body


# Binary Operations
class BinaryOperation(Expression):
    """Binary operation expression"""
    
    def __init__(self, operator: str, left: Expression, right: Expression):
        super().__init__("binary")
        self.operator = operator
        self.left = left
        self.right = right


class ChainOperation(Expression):
    """Chain operation (then)"""
    
    def __init__(self, left: Expression, right: Expression):
        super().__init__("chain")
        self.left = left
        self.right = right


# Utility functions for backward compatibility with tuple-based AST
def ast_to_tuple(node: ASTNode) -> tuple:
    """Convert AST node to tuple format for backward compatibility"""
    if isinstance(node, Program):
        return ("program", [ast_to_tuple(stmt) for stmt in node.statements])
    elif isinstance(node, Assignment):
        return ("assign", node.name, ast_to_tuple(node.value))
    elif isinstance(node, TupleAssignment):
        return ("tuple_assign", node.names, ast_to_tuple(node.value))
    elif isinstance(node, StringLiteral):
        return ("string", node.value)
    elif isinstance(node, NumberLiteral):
        return ("number", node.value)
    elif isinstance(node, BooleanLiteral):
        return ("bool", node.value)
    elif isinstance(node, ArrayLiteral):
        return ("array", [ast_to_tuple(elem) for elem in node.elements])
    elif isinstance(node, ObjectLiteral):
        return ("object", [(key, ast_to_tuple(value)) for key, value in node.pairs])
    elif isinstance(node, Identifier):
        return node.name
    elif isinstance(node, ToolDefinition):
        return ("tool", node.name, node.params, node.return_name, ast_to_tuple(node.body))
    elif isinstance(node, ToolCall):
        return ("tool_call", node.name, [ast_to_tuple(arg) for arg in node.args])
    elif isinstance(node, AgentDefinition):
        return ("agent_def", node.name, node.model, node.tools)
    elif isinstance(node, AgentCommand):
        return ("agent_command", ast_to_tuple(node.agent_ref), node.command, 
                ast_to_tuple(node.data) if node.data else None)
    elif isinstance(node, AgentBlock):
        return ("agent_block", ast_to_tuple(node.agent_ref), 
                [ast_to_tuple(stmt) for stmt in node.statements])
    elif isinstance(node, AgentReference):
        if node.ref_type == "alloy":
            return ("alloy",)
        elif node.ref_type == "agent_var":
            return ("agent_var", node.value)
        else:
            return ("agent", node.value)
    elif isinstance(node, AgentStatement):
        return (node.node_type, node.content if isinstance(node.content, str) else ast_to_tuple(node.content))
    elif isinstance(node, Block):
        return ("block", [ast_to_tuple(stmt) for stmt in node.statements])
    elif isinstance(node, WhenStatement):
        return ("when", ast_to_tuple(node.condition), ast_to_tuple(node.then_stmt), ast_to_tuple(node.otherwise_stmt))
    elif isinstance(node, ForLoop):
        return ("for", node.variable, ast_to_tuple(node.iterable), ast_to_tuple(node.body))
    elif isinstance(node, WhileLoop):
        return ("while", ast_to_tuple(node.condition), ast_to_tuple(node.body))
    elif isinstance(node, BinaryOperation):
        return ("binary", node.operator, ast_to_tuple(node.left), ast_to_tuple(node.right))
    elif isinstance(node, ChainOperation):
        return ("chain", ast_to_tuple(node.left), ast_to_tuple(node.right))
    else:
        # Fallback for any missed cases
        return (node.node_type,)