# neurotron/lang/frontend/parser/parser.py

from enum import Enum, auto
from neurotron.lang.frontend.lexer.lexer import Lexer
from neurotron.lang.frontend.tokens import TokenType
from neurotron.lang.frontend.AST import (
    Program, FunctionDef, IfExpr,
    BinOp, Call, Var, IntLiteral
)


class ParserError(Exception):
    pass


class Parser:
    def __init__(self, source: str):
        self.lexer = Lexer(source)
        self.tokens = list(self.lexer.tokenize())
        self.pos = 0

    def skip_newlines(self):
        while self.match(TokenType.NEWLINE):
            pass

    # -----------------------------------------------------
    # Basic token utilities
    # -----------------------------------------------------

    def peek(self):
        if self.pos < len(self.tokens):
            return self.tokens[self.pos]
        return None

    def advance(self):
        tok = self.peek()
        self.pos += 1
        return tok

    def match(self, *types):
        tok = self.peek()
        if tok and tok.type in types:
            return self.advance()
        return None

    def expect(self, type_: TokenType):
        tok = self.peek()
        if not tok or tok.type != type_:
            raise ParserError(f"Expected {type_}, got {tok}")
        return self.advance()

    # -----------------------------------------------------
    # Entry point
    # -----------------------------------------------------

    def parse(self) -> Program:
        functions = []

        # parse function definitions
        while self._is_start_of_function():
            functions.append(self.parse_function())

        # final expression
        self.skip_newlines()
        expr = self.parse_expr()
        self.skip_newlines()
        
        # 🔒 garantir fim de ficheiro
        self.expect(TokenType.EOF)

        return Program(functions, expr)

    def _is_start_of_function(self):
        tok = self.peek()
        return tok and tok.type == TokenType.DEF

    # -----------------------------------------------------
    # Function definition
    # -----------------------------------------------------

    def parse_function(self) -> FunctionDef:
        self.expect(TokenType.DEF)

        name_tok = self.expect(TokenType.IDENT)
        name = name_tok.value

        self.expect(TokenType.LPAREN)

        # only 0 or 1 argument for NL v0.1
        params = []
        if self.match(TokenType.IDENT):
            params.append(self.tokens[self.pos - 1].value)

        self.expect(TokenType.RPAREN)

        # newline + INDENT expected
        self.expect(TokenType.NEWLINE)
        self.expect(TokenType.INDENT)

        self.skip_newlines()
        body_expr = self.parse_expr()
        self.skip_newlines()

        # block ends with DEDENT
        self.expect(TokenType.DEDENT)

        return FunctionDef(name, params, body_expr)

    # -----------------------------------------------------
    # Expressions (recursive descent)
    # -----------------------------------------------------

    def parse_expr(self):
        """
        Handle if-expression:
            se cond entao expr senao expr
        """
        tok = self.peek()
        if tok and tok.type == TokenType.SE:
            return self.parse_if_expr()

        return self.parse_binary_expr()

    # -----------------------------------------------------
    # If-expression
    # -----------------------------------------------------

    def parse_if_expr(self):
        self.expect(TokenType.SE)
        cond = self.parse_binary_expr()

        # entao
        self.expect(TokenType.ENTAO)
        self.expect(TokenType.NEWLINE)
        self.expect(TokenType.INDENT)

        self.skip_newlines()
        then_expr = self.parse_expr()
        self.skip_newlines()

        self.expect(TokenType.DEDENT)

        # senao
        self.expect(TokenType.SENAO)
        self.expect(TokenType.NEWLINE)
        self.expect(TokenType.INDENT)

        self.skip_newlines()
        else_expr = self.parse_expr()
        self.skip_newlines()

        self.expect(TokenType.DEDENT)

        return IfExpr(cond, then_expr, else_expr)

    # -----------------------------------------------------
    # Binary expressions: < , + , -
    # Precedence:   < lowest
    #               + -
    # -----------------------------------------------------

    def parse_binary_expr(self):
        self.skip_newlines()              # 👈 NOVO
        left = self.parse_term()
        self.skip_newlines()              # 👈 NOVO

        tok = self.peek()
        while tok and tok.type in (TokenType.LT, TokenType.PLUS, TokenType.MINUS):
            op = tok.type
            self.advance()
            self.skip_newlines()          # 👈 NOVO
            right = self.parse_term()
            self.skip_newlines()          # 👈 NOVO
            left = BinOp(left, op, right)
            tok = self.peek()

        return left

    # -----------------------------------------------------
    # Term: literal, variable, call, parenthesis
    # -----------------------------------------------------

    def parse_term(self):
        self.skip_newlines()              # 👈 NOVO
        tok = self.peek()

        if tok.type == TokenType.INT:
            self.advance()
            return IntLiteral(tok.value)

        if tok.type == TokenType.IDENT:
            return self.parse_ident_or_call()

        if tok.type == TokenType.LPAREN:
            self.advance()
            expr = self.parse_expr()
            self.expect(TokenType.RPAREN)
            return expr

        raise ParserError(f"Unexpected token in term: {tok}")


    def parse_ident_or_call(self):
        name_tok = self.expect(TokenType.IDENT)
        name = name_tok.value

        # call: name(expr)
        if self.match(TokenType.LPAREN):
            args = []

            if self.peek().type != TokenType.RPAREN:
                args.append(self.parse_expr())

                # NL v0.1 supports only one argument anyway,
                # but we accept comma-less single arg.
                # Future: parse multiple args.

            self.expect(TokenType.RPAREN)
            return Call(name, args)

        # variable
        return Var(name)


# ---------------------------------------------------------
# Convenience entry point
# ---------------------------------------------------------

def parse(source: str) -> Program:
    return Parser(source).parse()
