# neurotron/lang/holodeck/vm.py
"""
Holodeck VM v0.1 — stack machine minimalista para a linguagem Neuro.

Compatível com:
- neurotron/lang/ir.py  (ModuleIR, FunctionIR, Instruction, Op)

Semântica v0.1 (fib-world):
- Instruções:
  PUSH_CONST int
  LOAD_VAR name
  STORE_VAR name
  ADD / SUB / LT
  LABEL name
  JUMP label
  JUMP_IF_FALSE label   (consome cond)
  CALL func_name nargs  (consome nargs args, produz 1 valor)
  RET                   (consome 1 valor e retorna)

Notas:
- Labels são resolvidas no início de cada frame (mapa label->ip).
- VM assume que o IR foi validado antes (validate_ir.py).
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

from neurotron.lang.ir import ModuleIR, FunctionIR, Instruction, Op


class VMError(RuntimeError):
    pass


@dataclass
class Frame:
    """
    Frame de execução de uma função.

    func: FunctionIR
    locals: {param_name -> value}
    ip: índice da próxima instrução
    labels: {label_name -> instruction_index}
    """

    func: FunctionIR
    locals: Dict[str, Any]
    ip: int = 0
    labels: Dict[str, int] = field(default_factory=dict)


@dataclass
class VMState:
    module: ModuleIR
    value_stack: List[Any] = field(default_factory=list)
    call_stack: List[Frame] = field(default_factory=list)
    halted: bool = False
    last_result: Optional[Any] = None


class HolodeckVM:
    def __init__(self, module_ir: ModuleIR):
        self.state = VMState(module=module_ir)

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

    def run(self, entrypoint: str = "__main__", args: Optional[List[Any]] = None) -> Any:
        if args is None:
            args = []

        func = self._get_function(entrypoint)
        if func is None:
            raise VMError(f"Função de entrada '{entrypoint}' não encontrada no módulo IR.")

        if len(args) != len(func.params):
            raise VMError(
                f"Função '{entrypoint}' espera {len(func.params)} args, mas recebeu {len(args)}."
            )

        # Reset state
        self.state.value_stack = []
        self.state.call_stack = []
        self.state.halted = False
        self.state.last_result = None

        # Push frame inicial
        locals_map = {name: value for name, value in zip(func.params, args)}
        self.state.call_stack.append(self._make_frame(func, locals_map))

        # Loop principal
        while not self.state.halted and self.state.call_stack:
            self._step()

        return self.state.last_result

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _get_function(self, name: str) -> Optional[FunctionIR]:
        funcs = getattr(self.state.module, "functions", None)
        if isinstance(funcs, dict):
            return funcs.get(name)
        return None

    def _make_frame(self, func: FunctionIR, locals_map: Dict[str, Any]) -> Frame:
        labels = self._resolve_labels(func)
        return Frame(func=func, locals=locals_map, ip=0, labels=labels)

    @staticmethod
    def _resolve_labels(func: FunctionIR) -> Dict[str, int]:
        """
        Constrói um mapa label->ip (índice da instrução onde LABEL aparece).
        Semântica: saltar para label significa continuar a execução A PARTIR do LABEL.
        (O LABEL em si é um no-op.)
        """
        label_map: Dict[str, int] = {}
        for ip, instr in enumerate(func.instructions):
            if isinstance(instr, Instruction) and instr.op == Op.LABEL:
                if len(instr.args) == 1 and isinstance(instr.args[0], str):
                    label_map[instr.args[0]] = ip
        return label_map

    def _current_frame(self) -> Frame:
        if not self.state.call_stack:
            raise VMError("Call stack vazia — sem frame atual.")
        return self.state.call_stack[-1]

    def _push(self, v: Any) -> None:
        self.state.value_stack.append(v)

    def _pop(self) -> Any:
        if not self.state.value_stack:
            raise VMError("Stack underflow: pop em stack vazia.")
        return self.state.value_stack.pop()

    def _step(self) -> None:
        frame = self._current_frame()
        instrs: List[Instruction] = frame.func.instructions

        if frame.ip < 0 or frame.ip >= len(instrs):
            # Se sair fora, tratamos como retorno implícito None
            self._return_from_function(None)
            return

        instr = instrs[frame.ip]
        frame.ip += 1  # avanço otimista

        if not isinstance(instr, Instruction) or not isinstance(instr.op, Op):
            raise VMError(f"Instrução inválida em {frame.func.name} ip={frame.ip-1}: {instr!r}")

        op = instr.op
        args = instr.args

        # ---------------------------
        # Data / locals
        # ---------------------------
        if op == Op.PUSH_CONST:
            value = args[0]
            self._push(value)
            return

        if op == Op.LOAD_VAR:
            name = args[0]
            if name not in frame.locals:
                raise VMError(f"Variável '{name}' não definida no frame de {frame.func.name}.")
            self._push(frame.locals[name])
            return

        if op == Op.STORE_VAR:
            name = args[0]
            frame.locals[name] = self._pop()
            return

        # ---------------------------
        # Arithmetic / compare
        # ---------------------------
        if op == Op.ADD:
            b = self._pop()
            a = self._pop()
            self._push(a + b)
            return

        if op == Op.SUB:
            b = self._pop()
            a = self._pop()
            self._push(a - b)
            return

        if op == Op.LT:
            b = self._pop()
            a = self._pop()
            self._push(a < b)
            return

        # ---------------------------
        # Control flow
        # ---------------------------
        if op == Op.LABEL:
            # no-op
            return

        if op == Op.JUMP:
            label = args[0]
            target = frame.labels.get(label)
            if target is None:
                raise VMError(f"JUMP: label '{label}' não existe em {frame.func.name}.")
            frame.ip = target
            return

        if op == Op.JUMP_IF_FALSE:
            label = args[0]
            cond = self._pop()
            if not cond:
                target = frame.labels.get(label)
                if target is None:
                    raise VMError(f"JUMP_IF_FALSE: label '{label}' não existe em {frame.func.name}.")
                frame.ip = target
            return

        # ---------------------------
        # Calls / returns
        # ---------------------------
        if op == Op.CALL:
            func_name = args[0]
            nargs = args[1]

            func_ir = self._get_function(func_name)
            if func_ir is None:
                raise VMError(f"CALL: função '{func_name}' não encontrada no módulo.")

            if nargs != len(func_ir.params):
                raise VMError(
                    f"CALL: '{func_name}' espera {len(func_ir.params)} args, recebeu {nargs}."
                )

            if nargs > len(self.state.value_stack):
                raise VMError(
                    f"CALL: stack tem {len(self.state.value_stack)} valores, mas {nargs} args foram pedidos."
                )

            # args estão no topo da stack, em ordem de avaliação; pop e reverter
            call_args = [self._pop() for _ in range(nargs)]
            call_args.reverse()

            locals_map = {name: value for name, value in zip(func_ir.params, call_args)}
            self.state.call_stack.append(self._make_frame(func_ir, locals_map))
            return

        if op == Op.RET:
            value = self._pop() if self.state.value_stack else None
            self._return_from_function(value)
            return

        raise VMError(f"Opcode não suportado no Holodeck v0.1: {op}")

    def _return_from_function(self, value: Any) -> None:
        if not self.state.call_stack:
            self.state.halted = True
            self.state.last_result = value
            return

        self.state.call_stack.pop()

        if not self.state.call_stack:
            self.state.halted = True
            self.state.last_result = value
        else:
            # retorno vira valor na stack do caller
            self._push(value)
