# neurotron/lang/frontend/ir_builder.py
"""
Gerador de IR (AST → IR) para Neuro v0.1

Cobre apenas o subconjunto necessário para fib.nl:
- inteiros
- variáveis
- chamadas de função
- +, -
- <
- se ... entao ... senao ...
- def de funções
- expressão final como __main__
"""

from __future__ import annotations

from typing import List

from neurotron.lang.ir import ModuleIR, FunctionIR, Op
from neurotron.lang.frontend.tokens import TokenType
from neurotron.lang.frontend.AST import (
    Program,
    FunctionDef,
    IfExpr,
    BinOp,
    Call,
    Var,
    IntLiteral,
)


class IRBuildError(Exception):
    pass


class IRBuilder:
    def __init__(self):
        self.module = ModuleIR()
        self._label_counter = 0
        self.current_function: FunctionIR | None = None

    # --------------------------------------------------
    # API pública
    # --------------------------------------------------

    def build(self, program: Program) -> ModuleIR:
        """
        Constrói o IR completo de um Program AST.
        - Gera IR para cada função definida.
        - Cria função especial __main__/0 para a expressão final.
        """

        # 1) Funções definidas pelo utilizador
        for fn in program.functions:
            self._build_function(fn)

        # 2) Função de entrada __main__
        main_fn = FunctionDef(
            name="__main__",
            params=[],
            body=program.final_expr,
        )
        self._build_function(main_fn, is_entry=True)

        return self.module

    # --------------------------------------------------
    # Funções
    # --------------------------------------------------

    def _build_function(self, fn: FunctionDef, is_entry: bool = False):
        prev_fn = self.current_function

        ir_fn = FunctionIR(fn.name, fn.params)
        self.current_function = ir_fn

        # Compilar corpo da função (é uma expressão única no v0.1)
        self._compile_expr(fn.body)

        # Garantir RET no fim
        ir_fn.emit(Op.RET)

        # Registar função no módulo
        self.module.add_function(ir_fn)

        if is_entry:
            self.module.entrypoint = fn.name

        self.current_function = prev_fn

    # --------------------------------------------------
    # Expressões
    # --------------------------------------------------

    def _compile_expr(self, node):
        """
        Compila uma expressão AST, deixando o resultado
        no topo da stack da função atual.
        """
        if isinstance(node, IntLiteral):
            return self._compile_int_literal(node)
        elif isinstance(node, Var):
            return self._compile_var(node)
        elif isinstance(node, Call):
            return self._compile_call(node)
        elif isinstance(node, BinOp):
            return self._compile_binop(node)
        elif isinstance(node, IfExpr):
            return self._compile_if_expr(node)
        else:
            raise IRBuildError(f"Tipo de expressão ainda não suportado: {type(node).__name__}")

    # --------------------------------------------------
    # Literais / variáveis
    # --------------------------------------------------

    def _compile_int_literal(self, node: IntLiteral):
        self._fn_emit(Op.PUSH_CONST, node.value)

    def _compile_var(self, node: Var):
        # NL v0.1: apenas parâmetros de função existem como variáveis.
        # A VM irá cuidar de resolver o nome no frame.
        self._fn_emit(Op.LOAD_VAR, node.name)

    # --------------------------------------------------
    # Chamadas de função
    # --------------------------------------------------

    def _compile_call(self, node: Call):
        # Estratégia: avaliar argumentos da esquerda para a direita,
        # empurrando-os na stack, depois CALL func nargs.
        for arg in node.args:
            self._compile_expr(arg)
        self._fn_emit(Op.CALL, node.func_name, len(node.args))

    # --------------------------------------------------
    # Operações binárias (+, -, <)
    # --------------------------------------------------

    def _compile_binop(self, node: BinOp):
        # Primeiro avalia left, depois right → na stack ficam [left, right]
        self._compile_expr(node.left)
        self._compile_expr(node.right)

        op = node.op

        if op == TokenType.PLUS:
            self._fn_emit(Op.ADD)
        elif op == TokenType.MINUS:
            self._fn_emit(Op.SUB)
        elif op == TokenType.LT:
            self._fn_emit(Op.LT)
        else:
            raise IRBuildError(f"Operador binário ainda não suportado: {op}")

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

    def _compile_if_expr(self, node: IfExpr):
        """
        se cond entao expr1 senao expr2

        IR em stack machine:
            [cond]
            JUMP_IF_FALSE L_else
            [expr1]
            JUMP L_end
        L_else:
            [expr2]
        L_end:
            # resultado de ambas as branches fica na stack
        """
        else_label = self._new_label("L_else")
        end_label = self._new_label("L_end")

        # cond
        self._compile_expr(node.cond)
        self._fn_emit(Op.JUMP_IF_FALSE, else_label)

        # then
        self._compile_expr(node.then_expr)
        self._fn_emit(Op.JUMP, end_label)

        # else:
        self._fn_emit(Op.LABEL, else_label)
        self._compile_expr(node.else_expr)

        # fim:
        self._fn_emit(Op.LABEL, end_label)

    # --------------------------------------------------
    # Utilitários internos
    # --------------------------------------------------

    def _new_label(self, prefix: str = "L") -> str:
        name = f"{prefix}{self._label_counter}"
        self._label_counter += 1
        return name

    def _fn_emit(self, op: Op, *args):
        if not self.current_function:
            raise IRBuildError("Nenhuma função corrente ativa ao emitir instrução")
        self.current_function.emit(op, *args)


# ------------------------------------------------------
# Conveniência: função de topo
# ------------------------------------------------------

def build_ir(program: Program) -> ModuleIR:
    """
    Função de conveniência: recebe AST Program, devolve ModuleIR.
    """
    builder = IRBuilder()
    return builder.build(program)
