# alloy_mvp_interpreter.py
import ply.lex as lex
import ply.yacc as yacc
from typing import Any, Dict, List, Tuple, Optional
import openai  # Or anthropic; we'll use OpenAI for demo
import os
from dotenv import load_dotenv

# --- Lexer ---

tokens = (
    'IDENTIFIER', 'STRING', 'NUMBER', 'BOOLEAN', 'LBRACKET', 'RBRACKET', 'LBRACE', 'RBRACE',
    'LPAREN', 'RPAREN', 'EQ', 'PLUS', 'MINUS', 'GT', 'GE', 'COMMA', 'COLON', 'TOOL', 'TAKES',
    'RETURNS', 'WHEN', 'THEN', 'OTHERWISE', 'FOR', 'IN', 'DO', 'WHILE', 'AGENT', 'ALLOY',
    'MUST', 'SHOULD', 'DATA', 'WITH', 'TOOLS'
)

# Reserved keywords
reserved = {
    'tool': 'TOOL', 'takes': 'TAKES', 'returns': 'RETURNS', 'when': 'WHEN', 'then': 'THEN',
    'otherwise': 'OTHERWISE', 'for': 'FOR', 'in': 'IN', 'do': 'DO', 'while': 'WHILE',
    'agent': 'AGENT', 'alloy': 'ALLOY', 'must': 'MUST', 'should': 'SHOULD', 'data': 'DATA',
    'with': 'WITH', 'tools': 'TOOLS',
    'true': 'BOOLEAN', 'false': 'BOOLEAN', 'null': 'NULL'
}

t_STRING = r'\"([^\\\n]|(\\.))*?\"|\'([^\\\n]|(\\.))*?\''
t_LBRACKET = r'\['
t_RBRACKET = r'\]'
t_LBRACE = r'\{'
t_RBRACE = r'\}'
t_LPAREN = r'\('
t_RPAREN = r'\)'
t_EQ = r'='
t_PLUS = r'\+'
t_MINUS = r'\-'
t_GT = r'>'
t_GE = r'>='
t_COMMA = r','
t_COLON = r':'

def t_IDENTIFIER(t):
    r'[a-zA-Z_][a-zA-Z_0-9]*'
    t.type = reserved.get(t.value, 'IDENTIFIER')
    return t

def t_NUMBER(t):
    r'\d+(\.\d+)?'
    t.value = float(t.value) if '.' in t.value else int(t.value)
    return t

t_ignore = ' \t\n'

def t_error(t):
    print(f"Illegal character '{t.value[0]}'")
    t.lexer.skip(1)

lexer = lex.lex()

# --- Parser ---

# Simplified grammar based on MVP
def p_program(p):
    '''program : statement_list'''
    p[0] = ('program', p[1])

def p_statement_list(p):
    '''statement_list : statement statement_list
                      | empty'''
    if len(p) == 3:
        p[0] = [p[1]] + p[2]
    else:
        p[0] = []

def p_statement(p):
    '''statement : assignment
                 | tuple_assignment
                 | agent_definition
                 | tool_def
                 | control_flow
                 | expression'''
    p[0] = p[1]

def p_assignment(p):
    '''assignment : IDENTIFIER EQ expression'''
    p[0] = ('assign', p[1], p[3])

def p_tuple_assignment(p):
    '''tuple_assignment : identifier_list EQ expression'''
    p[0] = ('tuple_assign', p[1], p[3])

def p_identifier_list(p):
    '''identifier_list : IDENTIFIER COMMA identifier_list
                       | IDENTIFIER'''
    if len(p) == 4:
        p[0] = [p[1]] + p[3]
    else:
        p[0] = [p[1]]

def p_agent_definition(p):
    '''agent_definition : IDENTIFIER EQ AGENT STRING
                        | IDENTIFIER EQ AGENT STRING WITH TOOLS tool_list'''
    if len(p) == 5:
        # Simple agent definition: claude = agent "claude-3.5-sonnet"
        p[0] = ('agent_def', p[1], p[4], [])
    else:
        # Agent with tools: researcher = agent "claude-3.5-sonnet" with tools web_search, read_files
        p[0] = ('agent_def', p[1], p[4], p[7])

def p_tool_list(p):
    '''tool_list : IDENTIFIER COMMA tool_list
                 | IDENTIFIER'''
    if len(p) == 4:
        p[0] = [p[1]] + p[3]
    else:
        p[0] = [p[1]]

def p_tool_def(p):
    '''tool_def : TOOL IDENTIFIER TAKES param_list RETURNS IDENTIFIER block'''
    p[0] = ('tool', p[2], p[4], p[6], p[7])

def p_param_list(p):
    '''param_list : IDENTIFIER COMMA param_list
                  | IDENTIFIER
                  | empty'''
    if len(p) == 4:
        p[0] = [p[1]] + p[3]
    elif len(p) == 2:
        p[0] = [p[1]]
    else:
        p[0] = []

def p_block(p):
    '''block : LBRACE statement_list RBRACE'''
    p[0] = ('block', p[2])

def p_control_flow_when(p):
    '''control_flow : WHEN expression THEN statement_or_block OTHERWISE statement_or_block'''
    p[0] = ('when', p[2], p[4], p[6])

def p_control_flow_for(p):
    '''control_flow : FOR IDENTIFIER IN expression DO statement_or_block'''
    p[0] = ('for', p[2], p[4], p[6])

def p_control_flow_while(p):
    '''control_flow : WHILE expression DO statement_or_block'''
    p[0] = ('while', p[2], p[4])

def p_statement_or_block(p):
    '''statement_or_block : statement
                          | block'''
    p[0] = p[1]

def p_expression(p):
    '''expression : literal
                  | IDENTIFIER
                  | tool_call
                  | agent_command
                  | agent_block
                  | binary_op
                  | chain_op'''
    p[0] = p[1]

def p_literal_string(p):
    '''literal : STRING'''
    p[0] = ('string', p[1].strip('"\''))

def p_literal_number(p):
    '''literal : NUMBER'''
    p[0] = ('number', p[1])

def p_literal_bool(p):
    '''literal : BOOLEAN'''
    p[0] = ('bool', p[1] == 'true')

def p_literal_array(p):
    '''literal : LBRACKET expression_list RBRACKET'''
    p[0] = ('array', p[2])

def p_literal_object(p):
    '''literal : LBRACE key_value_list RBRACE'''
    p[0] = ('object', dict(p[2]))

def p_expression_list(p):
    '''expression_list : expression COMMA expression_list
                       | expression
                       | empty'''
    if len(p) == 4:
        p[0] = [p[1]] + p[3]
    elif len(p) == 2:
        p[0] = [p[1]]
    else:
        p[0] = []

def p_key_value_list(p):
    '''key_value_list : IDENTIFIER COLON expression COMMA key_value_list
                      | IDENTIFIER COLON expression
                      | empty'''
    if len(p) == 6:
        p[0] = [(p[1], p[3])] + p[5]
    elif len(p) == 4:
        p[0] = [(p[1], p[3])]
    else:
        p[0] = []

def p_tool_call(p):
    '''tool_call : IDENTIFIER LPAREN expression_list RPAREN'''
    p[0] = ('tool_call', p[1], p[3])

def p_agent_command(p):
    '''agent_command : agent_ref STRING
                     | agent_ref IDENTIFIER expression
                     | agent_ref IDENTIFIER'''
    if len(p) == 3:
        # Two cases: string command or identifier command without data
        if isinstance(p[2], str) and (p[2].startswith('"') or p[2].startswith("'")):
            # Simple string command: claude "analyze this"
            p[0] = ('agent_command', p[1], p[2], None)
        else:
            # Command without data: claude Analyze
            p[0] = ('agent_command', p[1], p[2], None)
    elif len(p) == 4:
        # Command with data: claude Analyze user_data
        p[0] = ('agent_command', p[1], p[2], p[3])

def p_agent_block(p):
    '''agent_block : agent_ref LBRACE agent_statement_list RBRACE'''
    p[0] = ('agent_block', p[1], p[3])

def p_agent_ref(p):
    '''agent_ref : AGENT STRING
                 | ALLOY
                 | IDENTIFIER'''
    if len(p) == 2:
        if p[1] == 'alloy':
            p[0] = ('alloy')
        else:
            # Reference to an agent variable: claude, researcher, etc.
            p[0] = ('agent_var', p[1])
    else:
        # Old syntax: agent "model-name"
        p[0] = ('agent', p[2])

def p_agent_statement_list(p):
    '''agent_statement_list : agent_statement agent_statement_list
                            | empty'''
    if len(p) == 3:
        p[0] = [p[1]] + p[2]
    else:
        p[0] = []

def p_agent_statement(p):
    '''agent_statement : MUST STRING
                       | SHOULD STRING
                       | DATA COLON expression'''
    if p[1] == 'data':
        p[0] = ('data', p[3])
    else:
        p[0] = (p[1], p[2])

def p_binary_op(p):
    '''binary_op : expression PLUS expression
                 | expression GT expression
                 | expression GE expression'''
    op = {'+': 'add', '>': 'gt', '>=': 'ge'}[p[2]]
    p[0] = ('binary', op, p[1], p[3])

def p_chain_op(p):
    '''chain_op : expression THEN expression'''
    p[0] = ('chain', p[1], p[3])

def p_empty(p):
    'empty :'
    pass

def p_error(p):
    print(f"Syntax error at '{p.value}'")

parser = yacc.yacc()

# --- Interpreter ---

class RuntimeError(Exception):
    pass

class Environment:
    def __init__(self, parent: Optional['Environment'] = None):
        self.vars: Dict[str, Any] = {}
        self.tools: Dict[str, callable] = {}
        self.agents: Dict[str, Any] = {}  # Agent model names or 'alloy'
        self.parent = parent

    def define(self, name: str, value: Any):
        self.vars[name] = value

    def get(self, name: str) -> Any:
        if name in self.vars:
            return self.vars[name]
        if self.parent:
            return self.parent.get(name)
        raise RuntimeError(f"Undefined variable '{name}'")

    def define_tool(self, name: str, func: callable):
        self.tools[name] = func

    def get_tool(self, name: str) -> callable:
        if name in self.tools:
            return self.tools[name]
        if self.parent:
            return self.parent.get_tool(name)
        raise RuntimeError(f"Undefined tool '{name}'")

    def define_agent(self, name: str, model: str, tools: List[str] = None):
        self.agents[name] = {'model': model, 'tools': tools or []}

    def get_agent(self, name: str) -> Dict[str, Any]:
        if name in self.agents:
            return self.agents[name]
        if self.parent:
            return self.parent.get_agent(name)
        raise RuntimeError(f"Undefined agent '{name}'")

class Interpreter:
    def __init__(self):
        self.global_env = Environment()
        
        # Load environment variables from .env file
        load_dotenv()
        
        # Get OpenAI API key from environment
        api_key = os.getenv('OPENAI_API_KEY')
        if api_key:
            self.openai_client = openai.OpenAI(api_key=api_key)
        else:
            self.openai_client = None
            print("Warning: OPENAI_API_KEY not found in environment. Agent blocks will not work.")
            print("Create a .env file with: OPENAI_API_KEY=your_key_here")

    def interpret(self, ast):
        return self.visit(ast, self.global_env)

    def visit(self, node, env: Environment):
        # Handle the case where node is just a string (identifier)
        if isinstance(node, str):
            if node == 'return':
                return None  # Handle return keyword
            return env.get(node)
        
        if not isinstance(node, tuple):
            return node
            
        node_type = node[0]
        if node_type == 'program':
            for stmt in node[1]:
                self.visit(stmt, env)
        elif node_type == 'assign':
            value = self.visit(node[2], env)
            env.define(node[1], value)
            return value
        elif node_type == 'tuple_assign':
            identifiers = node[1]
            value = self.visit(node[2], env)
            # Handle tuple assignment for result, error patterns
            if isinstance(value, tuple) and len(value) == len(identifiers):
                for i, ident in enumerate(identifiers):
                    env.define(ident, value[i])
                return value
            else:
                # If not a tuple or wrong length, assign None to error variables
                for i, ident in enumerate(identifiers):
                    if i == 0:
                        env.define(ident, value)
                    else:
                        env.define(ident, None)
                return value
        elif node_type == 'agent_def':
            name, model, tools = node[1], node[2].strip('"\''), node[3]
            env.define_agent(name, model, tools)
            # Also define as a regular variable for reference
            env.define(name, ('agent_ref', name))
            return ('agent_ref', name)
        elif node_type == 'tool':
            params, ret_name, body = node[2], node[3], node[4]
            def tool_func(*args):
                tool_env = Environment(env)
                for param, arg in zip(params, args):
                    tool_env.define(param, arg)
                result = self.visit(body, tool_env)
                return result, None  # (result, error)
            env.define_tool(node[1], tool_func)
        elif node_type == 'when':
            cond = self.visit(node[1], env)
            if cond:
                return self.visit(node[2], env)
            else:
                return self.visit(node[3], env)
        elif node_type == 'for':
            iterable = self.visit(node[2], env)
            results = []
            for item in iterable:
                loop_env = Environment(env)
                loop_env.define(node[1], item)
                results.append(self.visit(node[3], loop_env))
            return results
        elif node_type == 'while':
            results = []
            while self.visit(node[1], env):
                results.append(self.visit(node[2], env))
            return results
        elif node_type == 'block':
            block_env = Environment(env)
            result = None
            for stmt in node[1]:
                if isinstance(stmt, str) and stmt == 'return':
                    continue  # Skip 'return' keyword
                result = self.visit(stmt, block_env)
            return result
        elif node_type == 'string':
            return node[1]
        elif node_type == 'number':
            return node[1]
        elif node_type == 'bool':
            return node[1]
        elif node_type == 'array':
            return [self.visit(e, env) for e in node[1]]
        elif node_type == 'object':
            return {k: self.visit(v, env) for k, v in node[1]}
        elif node_type == 'tool_call':
            tool = env.get_tool(node[1])
            args = [self.visit(arg, env) for arg in node[2]]
            return tool(*args)
        elif node_type == 'agent_command' or node_type == 'agent_block':
            return self.execute_agent(node, env)
        elif node_type == 'binary':
            left = self.visit(node[2], env)
            right = self.visit(node[3], env)
            if node[1] == 'add':
                return left + right
            elif node[1] == 'gt':
                return left > right
            elif node[1] == 'ge':
                return left >= right
        elif node_type == 'chain':
            res1, err1 = self.visit(node[1], env)
            if err1:
                return None, err1
            res2, err2 = self.visit(node[2], env)  # Assume chain passes res1 to next
            return res2, err2
        # TODO: Add more node types (e.g., agent def)

    def execute_agent(self, node, env: Environment) -> Tuple[Any, Optional[str]]:
        agent_ref = node[1]
        
        # Determine the agent model and tools
        if agent_ref[0] == 'alloy':
            model = 'gpt-4.1'  # Simulate alloy as GPT-4-1 for meta
            tools = []
        elif agent_ref[0] == 'agent_var':
            # Reference to an agent variable (clean syntax)
            agent_name = agent_ref[1]
            agent_info = env.get_agent(agent_name)
            model = agent_info['model']
            tools = agent_info['tools']
        else:
            # Old syntax: agent "model-name"
            model = agent_ref[1]
            tools = []
        
        if node[0] == 'agent_command':
            command = node[2]
            data = node[3]
            
            # Check if command is capitalized (Command) vs lowercase (tool)
            if isinstance(command, str):
                if command[0].isupper():
                    # This is a Command
                    prompt = f"Execute command: {command}"
                    if data:
                        data_value = self.visit(data, env) if data else None
                        prompt += f" with data: {data_value}"
                else:
                    # This is a tool call - check if agent has this tool
                    if command not in tools:
                        return None, f"Agent does not have tool '{command}'"
                    prompt = f"Use tool '{command}'"
                    if data:
                        data_value = self.visit(data, env) if data else None
                        prompt += f" with data: {data_value}"
            else:
                # String command
                prompt = f"Execute: {command.strip('\"\'')}"
        else:  # agent_block
            statements = node[2]
            musts, shoulds, data = [], [], None
            for stmt in statements:
                if stmt[0] == 'must':
                    musts.append(stmt[1])
                elif stmt[0] == 'should':
                    shoulds.append(stmt[1])
                elif stmt[0] == 'data':
                    data = self.visit(stmt[1], env)
            prompt = "Execute agent block:\nMust: " + "\n".join(musts) + "\nShould: " + "\n".join(shoulds)
            if data:
                prompt += f"\nData: {data}"
        
        if not self.openai_client:
            return None, "OpenAI API key not configured. Please set OPENAI_API_KEY in your .env file."
            
        try:
            response = self.openai_client.chat.completions.create(
                model=model,
                messages=[{"role": "system", "content": "You are an Alloy agent."}, {"role": "user", "content": prompt}]
            )
            return response.choices[0].message.content, None
        except Exception as e:
            return None, str(e)

# --- Example Usage ---

code = '''
claude = agent "claude-3.5-sonnet"
researcher = agent "claude-3.5-sonnet" with tools web_search, read_files

analysis, error = claude Analyze "user data"
summary, error = researcher Summarize analysis

tool add_numbers takes a, b returns sum {
    sum = a + b
    return sum
}

result = add_numbers(5, 3)
'''

ast = parser.parse(code)
interpreter = Interpreter()
interpreter.interpret(ast)
print(interpreter.global_env.get('result'))  # Should print (8, None) or similar