"""
Lexical Analyzer for the Alloy Programming Language

This module provides the lexer (tokenizer) for the Alloy language.
It uses PLY (Python Lex-Yacc) to define tokens and lexical rules.
"""

import ply.lex as lex
from typing import Any


class AlloyLexer:
    """
    Lexical analyzer for the Alloy programming language.
    
    Defines all tokens, keywords, and lexical rules for tokenizing Alloy source code.
    """
    
    # Token definitions - these need to be available at class level for PLY
    tokens = (
        'IDENTIFIER', 'STRING', 'NUMBER', 'BOOLEAN', 'NULL',
        'LBRACKET', 'RBRACKET', 'LBRACE', 'RBRACE',
        'LPAREN', 'RPAREN', 'NE', 'EQEQ', 'EQ', 'PLUS', 'MINUS', 'MULTIPLY',
        'GT', 'GE', 'LT', 'LE', 'COMMA', 'COLON',
        'TOOL', 'TAKES', 'RETURNS', 'WHEN', 'THEN', 'OTHERWISE',
        'FOR', 'IN', 'DO', 'WHILE', 'AGENT', 'ALLOY',
        'MUST', 'SHOULD', 'WITH', 'TOOLS', 'COMMAND', 'MEANS'
    )

    # Reserved keywords mapping
    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',
        'with': 'WITH',
        'tools': 'TOOLS',
        'command': 'COMMAND',
        'means': 'MEANS',
        'true': 'BOOLEAN',
        'false': 'BOOLEAN',
        'null': 'NULL'
    }

    # Token rules (simple tokens)
    t_STRING = r'\"([^\\\n]|(\\.))*?\"|\'([^\\\n]|(\\.))*?\''
    t_LBRACKET = r'\['
    t_RBRACKET = r'\]'
    t_LBRACE = r'\{'
    t_RBRACE = r'\}'
    t_LPAREN = r'\('
    t_RPAREN = r'\)'
    t_NE = r'!='
    t_EQEQ = r'=='
    t_EQ = r'='
    t_PLUS = r'\+'
    t_MINUS = r'\-'
    t_MULTIPLY = r'\*'
    t_GT = r'>'
    t_GE = r'>='
    t_LT = r'<'
    t_LE = r'<='
    t_COMMA = r','
    t_COLON = r':'

    # Ignored characters (spaces, tabs, newlines)
    t_ignore = ' \t\n'

    def __init__(self):
        """Initialize the lexer"""
        self.lexer = None
        self.build()

    def t_IDENTIFIER(self, t):
        r'[a-zA-Z_][a-zA-Z_0-9]*'
        # Check if it's a reserved keyword
        t.type = self.reserved.get(t.value, 'IDENTIFIER')
        return t

    def t_NUMBER(self, t):
        r'\d+(\.\d+)?'
        # Convert to appropriate numeric type
        if '.' in t.value:
            t.value = float(t.value)
        else:
            t.value = int(t.value)
        return t

    def t_newline(self, t):
        r'\n+'
        t.lexer.lineno += len(t.value)

    def t_error(self, t):
        """Handle lexical errors"""
        print(f"Illegal character '{t.value[0]}' at line {t.lineno}")
        t.lexer.skip(1)

    def t_COMMENT(self, t):
        r'\#.*'
        # Comments are ignored
        pass

    def build(self):
        """Build the lexer"""
        self.lexer = lex.lex(module=self)
        return self.lexer

    def tokenize(self, data: str):
        """
        Tokenize input data.
        
        Args:
            data: Input string to tokenize
            
        Yields:
            Tokens from the input
        """
        self.lexer.input(data)
        while True:
            tok = self.lexer.token()
            if not tok:
                break
            yield tok

    def test(self, data: str):
        """
        Test the lexer with input data and print tokens.
        
        Args:
            data: Input string to tokenize
        """
        print(f"Tokenizing: {data}")
        print("-" * 40)
        
        for token in self.tokenize(data):
            print(f"{token.type:12} {token.value}")


def create_lexer():
    """
    Factory function to create a new lexer instance.
    
    Returns:
        AlloyLexer instance
    """
    return AlloyLexer()


# Global lexer instance for backward compatibility
_global_lexer = None

def get_global_lexer():
    """
    Get the global lexer instance, creating it if necessary.
    
    Returns:
        PLY lexer object for backward compatibility
    """
    global _global_lexer
    if _global_lexer is None:
        _global_lexer = AlloyLexer()
    return _global_lexer.lexer


# For backward compatibility - expose the PLY lexer directly
lexer = get_global_lexer()


if __name__ == "__main__":
    # Test the lexer
    test_code = '''
    tool add_numbers takes a, b returns sum {
        sum = a + b
        return sum
    }
    
    claude = agent "gpt-4o-mini"
    result, error = claude "analyze this data"
    '''
    
    alloy_lexer = AlloyLexer()
    alloy_lexer.test(test_code)