"""
Nexa Language — Stack-based Bytecode VM (v22)

Executes NexaBytecode produced by compiler.py.
10–50× faster than the AST tree-walking interpreter.
"""
from __future__ import annotations
import json, operator
from typing import Any, Dict, List, Optional


# ── VM Exceptions ─────────────────────────────────────────────────────────────
class VMError(Exception): pass
class VMReturn(Exception):
    def __init__(self, value): self.value = value
class VMBreak(Exception): pass
class VMContinue(Exception): pass


# ── Frame ─────────────────────────────────────────────────────────────────────
class Frame:
    def __init__(self, code: list, locals_: dict, globals_: dict):
        self.code = code        # list of [op, arg] pairs
        self.ip = 0             # instruction pointer
        self.locals = locals_   # local variable scope
        self.globals = globals_ # module-level globals
        self.stack: List[Any] = []
        self.iter_stack: List[Any] = []   # iterator stack for for-loops

    def push(self, v): self.stack.append(v)
    def pop(self):     return self.stack.pop()
    def peek(self):    return self.stack[-1]

    def load(self, name) -> Any:
        if name in self.locals:  return self.locals[name]
        if name in self.globals: return self.globals[name]
        raise VMError(f"Undefined variable: '{name}'")

    def store(self, name, val): self.locals[name] = val


# ── VM ────────────────────────────────────────────────────────────────────────
class NexaVM:
    BINOPS = {
        "+": operator.add, "-": operator.sub,
        "*": operator.mul, "/": operator.truediv,
        "//": operator.floordiv, "%": operator.mod,
        "**": operator.pow,
        "==": operator.eq, "!=": operator.ne,
        "<": operator.lt, ">": operator.gt,
        "<=": operator.le, ">=": operator.ge,
        "and": lambda a, b: a and b,
        "or":  lambda a, b: a or b,
        "&": operator.and_, "|": operator.or_,
        "^": operator.xor, "<<": operator.lshift,
        ">>": operator.rshift,
    }

    def __init__(self):
        # Pull all builtins into global scope
        from nexa_lang.builtins import BUILTINS
        self._globals: Dict[str, Any] = {
            **BUILTINS,
            "print": lambda *a: print(*a),
            "len":   len,
            "str":   str,
            "int":   int,
            "float": float,
            "bool":  bool,
            "type":  type,
            "range": lambda *a: list(range(*[int(x) for x in a])),
            "append": lambda lst, val: lst + [val],
            "true": True, "false": False, "null": None,
        }
        self._class_registry: Dict[str, dict] = {}

    # ── Public API ────────────────────────────────────────────────────────────
    def execute(self, bytecode) -> Any:
        """Execute a NexaBytecode object."""
        frame = Frame(
            [list(i) if hasattr(i, '__iter__') else [i.op, i.arg]
             for i in bytecode.instructions],
            {}, self._globals
        )
        return self._run_frame(frame)

    def execute_dict(self, d: dict) -> Any:
        """Execute bytecode from a .nexac dict."""
        from nexa_lang.compiler import NexaBytecode
        bc = NexaBytecode.from_dict(d)
        return self.execute(bc)

    # ── Frame execution ───────────────────────────────────────────────────────
    def _run_frame(self, frame: Frame) -> Any:
        while frame.ip < len(frame.code):
            instr = frame.code[frame.ip]
            op = instr[0] if isinstance(instr, list) else instr.op
            arg = instr[1] if isinstance(instr, list) else instr.arg
            frame.ip += 1
            result = self._dispatch(op, arg, frame)
            if result is VMReturn:
                return frame.pop()
        return None

    def _dispatch(self, op: str, arg: Any, f: Frame):
        # ── Load / Store ──────────────────────────────────────────────────────
        if op == "LOAD_CONST":    f.push(arg)
        elif op == "LOAD_NAME":   f.push(f.load(arg))
        elif op == "STORE_NAME":  f.store(arg, f.pop())
        elif op == "LOAD_ATTR":
            obj = f.pop()
            val = obj.get(arg) if isinstance(obj, dict) else getattr(obj, arg, None)
            f.push(val)
        elif op == "STORE_ATTR":
            val = f.pop(); obj = f.pop()
            if isinstance(obj, dict): obj[arg] = val
            else: setattr(obj, arg, val)
        elif op == "LOAD_INDEX":
            idx = f.pop(); obj = f.pop()
            f.push(obj[idx])
        elif op == "STORE_INDEX":
            val = f.pop(); idx = f.pop(); obj = f.pop()
            obj[idx] = val
        # ── Stack ops ────────────────────────────────────────────────────────
        elif op == "POP":   f.pop()
        elif op == "DUP":   f.push(f.peek())
        elif op == "PRINT": print(f.pop())
        # ── Build ────────────────────────────────────────────────────────────
        elif op == "BUILD_LIST":
            n = arg; items = [f.pop() for _ in range(n)][::-1]; f.push(items)
        elif op == "BUILD_DICT":
            n = arg; d = {}
            pairs = [(f.pop(), f.pop()) for _ in range(n)][::-1]
            for k, v in pairs: d[k] = v
            f.push(d)
        elif op == "BUILD_RANGE":
            end = f.pop(); start = f.pop()
            f.push(list(range(int(start), int(end) + (1 if arg else 0))))
        elif op == "BUILD_FSTRING":
            result = ""
            expr_vals = []
            for kind, _ in reversed(arg):
                if kind == "expr": expr_vals.append(f.pop())
            for kind, lit in arg:
                if kind == "lit": result += str(lit)
                else: result += str(expr_vals.pop())
            f.push(result)
        elif op == "GENERIC_RUNTIME_CHECK":
            target = f.pop()
            if isinstance(target, list):
                type_map = {"Int": int, "Float": float, "String": str, "Bool": bool}
                expected = type_map.get(arg)
                if expected:
                    for i, item in enumerate(target):
                        if not isinstance(item, expected):
                            raise VMError(f"Generic type violation: element [{i}] expected {arg}, got {type(item).__name__}")
            f.push(target)
        elif op == "MATCH_BIND":
            subject = f.pop()
            if subject is None: f.push(False)
            else:
                f.store(arg, subject)
                f.push(True)
        # ── Arithmetic / Logic ────────────────────────────────────────────────
        elif op == "BINARY_OP":
            b = f.pop(); a = f.pop()
            fn = self.BINOPS.get(arg)
            if fn is None: raise VMError(f"Unknown binary op: {arg}")
            f.push(fn(a, b))
        elif op == "UNARY_OP":
            a = f.pop()
            if arg == "-": f.push(-a)
            elif arg in ("not", "!"): f.push(not a)
            elif arg == "~": f.push(~a)
        # ── Jumps ────────────────────────────────────────────────────────────
        elif op == "JUMP":          f.ip = arg
        elif op == "JUMP_IF_FALSE":
            if not f.pop(): f.ip = arg
        elif op == "JUMP_IF_TRUE":
            if f.pop(): f.ip = arg
        elif op == "JUMP_IF_NOT_EQ":
            b = f.pop(); a = f.pop()
            if a != b: f.ip = arg
        elif op == "JUMP_IF_NOT_NULL":
            if f.peek() is not None: f.ip = arg
        # ── Iteration ────────────────────────────────────────────────────────
        elif op == "ITER_START":
            iterable = f.pop()
            f.iter_stack.append(iter(iterable if isinstance(iterable, list) else [iterable]))
        elif op == "ITER_NEXT":
            try:
                val = next(f.iter_stack[-1])
                f.push(val)
            except StopIteration:
                f.iter_stack.pop()
                f.ip = arg  # jump past loop
        elif op == "ITER_END":
            if f.iter_stack: f.iter_stack.pop()
        # ── Exceptions ────────────────────────────────────────────────────────
        elif op == "RAISE":
            val = f.pop(); raise VMError(str(val))
        elif op == "TRY_BEGIN": pass  # handled by TRY_END
        elif op == "ASSERT_NOT_NULL":
            val = f.pop()
            if val is None: raise VMError("Non-null assertion failed: value is null")
            f.push(val)
        # ── Functions ────────────────────────────────────────────────────────
        elif op == "MAKE_FUNC":
            fn_def = arg
            def make_closure(fd, captured_globals):
                def fn(*args):
                    locs = dict(zip(fd["params"], args))
                    if fd.get("is_abstract"):
                        raise VMError(f"Abstract method '{fd['name']}' not implemented")
                    from nexa_lang.compiler import NexaBytecode
                    sub_bc = NexaBytecode.from_dict(fd["code"]) if fd.get("code") else None
                    if sub_bc is None: return None
                    sub_frame = Frame(
                        [[i.op, i.arg] for i in sub_bc.instructions],
                        locs, captured_globals
                    )
                    try:
                        return self._run_frame(sub_frame)
                    except VMReturn as r:
                        return r.value
                fn.__name__ = fd["name"]
                return fn
            f.push(make_closure(fn_def, {**f.globals, **f.locals}))
            # STORE_NAME follows immediately
        elif op == "CALL":
            nargs = arg["nargs"]; nkwargs = arg.get("nkwargs", 0)
            kwargs = {}
            for _ in range(nkwargs):
                k = f.pop(); v = f.pop(); kwargs[k] = v
            args = [f.pop() for _ in range(nargs)][::-1]
            callee = f.pop()
            try:
                result = callee(*args, **kwargs) if kwargs else callee(*args)
            except TypeError as e:
                raise VMError(f"Call error: {e}")
            f.push(result)
        elif op == "RETURN":
            return VMReturn
        # ── Classes ────────────────────────────────────────────────────────────
        elif op == "MAKE_CLASS":
            cls_def = arg
            cls = type(cls_def["name"], (object,), {
                "__class_name__": cls_def["name"],
                **{m: self._make_method(md) for m, md in cls_def["methods"].items()}
            })
            self._class_registry[cls_def["name"]] = cls
            f.push(cls)
        elif op == "IMPORT_NAME":
            from nexa_lang.builtins import BUILTINS
            mod = arg["module"]; alias = arg["alias"] or mod
            if mod in BUILTINS: f.locals[alias] = BUILTINS[mod]
        elif op == "HALT":
            return
        else:
            raise VMError(f"Unknown opcode: {op}")

    def _make_method(self, md: dict):
        def method(self_obj, *args):
            from nexa_lang.compiler import NexaBytecode
            sub_bc = NexaBytecode.from_dict(md["code"])
            locs = {"self": self_obj}
            locs.update(zip(md["params"], args))
            sub_frame = Frame([[i.op, i.arg] for i in sub_bc.instructions],
                              locs, self._globals)
            try:
                return self._run_frame(sub_frame)
            except VMReturn as r:
                return r.value
        return method
