# src/compiler/ast.py

@dataclass
class Expression:
    """Base class for AST nodes representing expressions."""

@dataclass
class Literal(Expression):
    value: int | bool | None
    # (value=None is used when parsing the keyword `unit`)

@dataclass
class Identifier(Expression):
    name: str

@dataclass
class BinaryOp(Expression):
    """AST node for a binary operation like `A + B`"""
    left: Expression
    op: str
    right: Expression

... # You get to define more later

# Construct AST for "x + 3" using
example_ast = BinaryOp(
    left=Identifier(name="x"),
    op="+",
    right=Literal(value=3),
)

# src/compiler/parser.py

from compiler.tokenizer import Token
import compiler.ast as ast

def parse(tokens: list[Token]) -> ast.Expression:
    ...

def parse(tokens: list[Token]) -> ast.Expression:
    # This keeps track of which token we're looking at.
    pos = 0

    # Let's first define two useful helper functions:
    # 'peek' and 'consume'.

    # 'peek()' returns the token at 'pos',
    # or a special 'end' token if we're past the end
    # of the token list.
    # This way we don't have to worry about going past
    # the end elsewhere.
    def peek() -> Token:
        if pos < len(tokens):
            return tokens[pos]
        else:
            return Token(
                location=tokens[-1].location,
                type="end",
                text="",
            )

    # 'consume(expected)' returns the token at 'pos'
    # and moves 'pos' forward.
    #
    # If the optional parameter 'expected' is given,
    # it checks that the token being consumed has that text.
    # If 'expected' is a list, then the token must have
    # one of the texts in the list.
    def consume(expected: str | list[str] | None = None) -> Token:
        token = peek()
        if isinstance(expected, str) and token.text != expected:
            raise Exception(f'{token.location}: expected "{expected}"')
        if isinstance(expected, list) and token.text not in expected:
            comma_separated = ", ".join([f'"{e}"' for e in expected])
            raise Exception(f'{token.location}: expected one of: {comma_separated}')
        pos += 1
        return token

    # This is the parsing function for integer literals.
    # It checks that we're looking at an integer literal token,
    # moves past it, and returns a 'Literal' AST node
    # containing the integer from the token.
    def parse_int_literal() -> ast.Literal:
        if peek().type != 'int_literal':
            raise Exception(f'{peek().location}: expected an integer literal')
        token = consume()
        return ast.Literal(int(token.text))

    # This is our main parsing function for this example.
    # To parse "integer + integer" expressions,
    # it uses `parse_int_literal` to parse the first integer,
    # then it checks that there's a supported operator,
    # and finally it uses `parse_int_literal` to parse the
    # second integer.
    def parse_expression() -> ast.BinaryOp:
        left = parse_int_literal()
        operator_token = consume(['+', '-'])
        right = parse_int_literal()
        return ast.BinaryOp(
            left,
            operator_token.text,
            right
        )

    return parse_expression()

def parse(tokens: list[Token]) -> ast.Expression:
    ...

    def parse_int_literal() -> ast.Literal:
        ... # (same as before)

    def parse_identifier() -> ast.Identifier
        ... # (very similar to parse_int_literal)

    # Borrowing from mathematics, we say that
    # "an identifier or an integer literal"
    # is called a "term".
    def parse_term() -> ast.Expression:
        if peek().type == 'int_literal':
            return parse_int_literal()
        elif peek().type == 'identifier':
            return parse_identifier()
        else:
            raise Exception(f'{peek().location}: expected an integer literal or an identifier')

    def parse_expression() -> ast.Expression:
        left = parse_term()
        operator_token = consume(['+', '-'])
        right = parse_term()
        return ast.BinaryOp(
            left,
            operator_token.text,
            right
        )

    return parse_expression()

1 - 2 + 3  =  1 - (2 + 3)  =  -4   😱

def parse_expression() -> ast.Expression:
    # Parse the first term.
    left = parse_term()

    # While there are more `+` or '-'...
    while peek().text in ['+', '-']:
        # Move past the '+' or '-'.
        operator_token = consume()
        operator = operator_token.text

        # Parse the operator on the right.
        right = parse_term()

        # Combine it with the stuff we've
        # accumulated on the left so far.
        left = ast.BinaryOp(
            left,
            operator,
            right
        )
    
    return left

def parse_expression() -> ast.Expression:
    # Same as before
    left = parse_term()
    while peek().text in ['+', '-']:
        operator_token = consume()
        operator = operator_token.text
        right = parse_term()
        left = ast.BinaryOp(
            left,
            operator,
            right
        )
    return left

def parse_term() -> ast.Expression:
    # Same structure as in 'parse_expression',
    # but the operators and function calls differ.
    left = parse_factor()
    while peek().text in ['*', '/']:
        operator_token = consume()
        operator = operator_token.text
        right = parse_factor()
        left = ast.BinaryOp(
            left,
            operator,
            right
        )
    return left

def parse_factor() -> ast.Expression:
    # Same as `parse_term` in the previous version.
    if peek().type == 'int_literal':
        return parse_int_literal()
    elif peek().type == 'identifier':
        return parse_identifier()
    else:
        raise Exception(f'{peek().location}: expected an integer literal or an identifier')

def parse_factor() -> ast.Expression:
    if peek().text == '(':
        return parse_parenthesized()
    elif peek().type == 'int_literal':
        return parse_int_literal()
    elif peek().type == 'identifier':
        return parse_identifier()
    else:
        raise Exception(f'{peek().location}: expected "(", an integer literal or an identifier')

def parse_parenthesized() -> ast.Expression:
    consume('(')
    # Recursively call the top level parsing function
    # to parse whatever is inside the parentheses.
    expr = parse_expression()
    consume(')')
    return expr