"""
Nexa Language — Bytecode Compiler (v22)
Compiles AST → flat instruction list for the stack-based VM.

Opcodes:
  LOAD_CONST, LOAD_NAME, STORE_NAME, LOAD_ATTR, STORE_ATTR
  LOAD_INDEX, STORE_INDEX, BUILD_LIST, BUILD_DICT, BUILD_SET
  BINARY_OP, UNARY_OP, CALL, RETURN, MAKE_FUNC, MAKE_CLASS
  JUMP, JUMP_IF_FALSE, JUMP_IF_TRUE
  POP, DUP, PRINT, HALT, IMPORT_NAME
  ITER_START, ITER_NEXT, ITER_END
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, List, Optional, Dict

from nexa_lang.ast_nodes import *


@dataclass
class Instruction:
    op: str
    arg: Any = None

    def __repr__(self):
        return f"{self.op:<18} {repr(self.arg) if self.arg is not None else ''}"


class NexaBytecode:
    """Container for compiled Nexa bytecode."""
    VERSION = "22.0"

    def __init__(self, source: str = "<unknown>"):
        self.source = source
        self.instructions: List[Instruction] = []
        self.constants: List[Any] = []
        self.names: List[str] = []

    def emit(self, op: str, arg: Any = None) -> int:
        """Emit instruction, return its index."""
        self.instructions.append(Instruction(op, arg))
        return len(self.instructions) - 1

    def patch(self, idx: int, arg: Any):
        """Patch a previously emitted instruction's arg (for forward jumps)."""
        self.instructions[idx].arg = arg

    def to_dict(self) -> dict:
        import hashlib, json
        code = [[i.op, i.arg] for i in self.instructions]
        raw = json.dumps(code)
        return {
            "nexa_version": self.VERSION,
            "source": self.source,
            "checksum": "sha256:" + hashlib.sha256(raw.encode()).hexdigest()[:12],
            "code": code,
        }

    @classmethod
    def from_dict(cls, d: dict) -> "NexaBytecode":
        bc = cls(d.get("source", "<unknown>"))
        bc.instructions = [Instruction(op, arg) for op, arg in d["code"]]
        return bc


class Compiler:
    """Compile a Nexa AST Program into NexaBytecode."""

    def __init__(self, source: str = "<unknown>"):
        self.bc = NexaBytecode(source)
        self._loop_stack: List[List[int]] = []   # list of break-patch indexes per loop

    def compile(self, program: Program) -> NexaBytecode:
        for stmt in program.statements:
            self._compile_node(stmt)
        self.bc.emit("HALT")
        return self.bc

    # ── Emit helpers ──────────────────────────────────────────────────────────
    def _emit(self, op, arg=None) -> int:
        return self.bc.emit(op, arg)

    def _patch(self, idx, arg):
        self.bc.patch(idx, arg)

    # ── Dispatch ──────────────────────────────────────────────────────────────
    def _compile_node(self, node):
        t = type(node).__name__
        method = getattr(self, f"_c_{t}", self._c_fallback)
        method(node)

    def _c_fallback(self, node):
        # Expressions used as statements — compile and pop
        self._compile_expr(node)
        self._emit("POP")

    def _compile_expr(self, node):
        t = type(node).__name__
        method = getattr(self, f"_e_{t}", self._e_fallback)
        method(node)

    def _e_fallback(self, node):
        self._emit("LOAD_CONST", repr(node))

    # ── Statement compilers ───────────────────────────────────────────────────
    def _c_Program(self, node: Program):
        for stmt in node.statements:
            self._compile_node(stmt)

    def _c_LetStatement(self, node):
        self._compile_expr(node.value)
        self._emit("STORE_NAME", node.name)

    def _c_Assignment(self, node):
        self._compile_expr(node.value)
        self._emit("STORE_NAME", node.name)

    def _c_MemberAssignment(self, node):
        self._compile_expr(node.obj)
        self._compile_expr(node.value)
        self._emit("STORE_ATTR", node.member)

    def _c_IndexAssignment(self, node):
        self._compile_expr(node.obj)
        self._compile_expr(node.index)
        self._compile_expr(node.value)
        self._emit("STORE_INDEX")

    def _c_ExprStatement(self, node):
        self._compile_expr(node.expr)
        self._emit("POP")

    def _c_ReturnStatement(self, node):
        if node.value:
            self._compile_expr(node.value)
        else:
            self._emit("LOAD_CONST", None)
        self._emit("RETURN")

    def _e_IfExpr(self, node):
        self._compile_expr(node.condition)
        jump_false = self._emit("JUMP_IF_FALSE", None)
        for stmt in node.then_block:
            self._compile_node(stmt)
        if node.else_block:
            jump_end = self._emit("JUMP", None)
            self._patch(jump_false, len(self.bc.instructions))
            for stmt in node.else_block:
                self._compile_node(stmt)
            self._patch(jump_end, len(self.bc.instructions))
        else:
            self._patch(jump_false, len(self.bc.instructions))

    def _c_WhileStatement(self, node):
        loop_start = len(self.bc.instructions)
        break_patches = []
        self._loop_stack.append(break_patches)
        self._compile_expr(node.condition)
        jump_end = self._emit("JUMP_IF_FALSE", None)
        for stmt in node.body:
            self._compile_node(stmt)
        self._emit("JUMP", loop_start)
        end = len(self.bc.instructions)
        self._patch(jump_end, end)
        for bp in break_patches:
            self._patch(bp, end)
        self._loop_stack.pop()

    def _c_LoopStatement(self, node):
        loop_start = len(self.bc.instructions)
        break_patches = []
        self._loop_stack.append(break_patches)
        for stmt in node.body:
            self._compile_node(stmt)
        self._emit("JUMP", loop_start)
        end = len(self.bc.instructions)
        for bp in break_patches:
            self._patch(bp, end)
        self._loop_stack.pop()

    def _c_ForStatement(self, node):
        self._compile_expr(node.iterable)
        self._emit("ITER_START")
        loop_top = len(self.bc.instructions)
        iter_next = self._emit("ITER_NEXT", None)  # jumps past loop if exhausted
        break_patches = [iter_next]
        self._loop_stack.append(break_patches)
        self._emit("STORE_NAME", node.var)
        for stmt in node.body:
            self._compile_node(stmt)
        self._emit("JUMP", loop_top)
        end = len(self.bc.instructions)
        self._patch(iter_next, end)
        for bp in break_patches:
            self._patch(bp, end)
        self._loop_stack.pop()

    def _c_BreakStatement(self, node):
        if self._loop_stack:
            bp = self._emit("JUMP", None)
            self._loop_stack[-1].append(bp)

    def _c_ContinueStatement(self, node):
        # Jump back to the top of the current loop
        if self._loop_stack:
            loop_top = self._find_loop_top()
            self._emit("JUMP", loop_top)

    def _find_loop_top(self) -> int:
        # Walk back to find the ITER_NEXT or last JUMP_IF_FALSE — estimate
        for i in range(len(self.bc.instructions)-1, -1, -1):
            if self.bc.instructions[i].op in ("ITER_NEXT", "JUMP_IF_FALSE"):
                return i
        return 0

    def _c_FunctionDef(self, node):
        # Compile function body into a sub-bytecode
        sub = Compiler(self.bc.source)
        sub_bc = sub.compile(Program(node.body + [ReturnStatement(NullLiteral())]))
        self._emit("MAKE_FUNC", {
            "name": node.name,
            "params": node.params,
            "code": sub_bc.to_dict(),
            "is_async": node.is_async,
        })
        self._emit("STORE_NAME", node.name)

    def _c_ClassDef(self, node):
        methods = {}
        for item in node.body:
            if isinstance(item, FunctionDef):
                sub = Compiler(self.bc.source)
                sub_bc = sub.compile(Program(item.body + [ReturnStatement(NullLiteral())]))
                methods[item.name] = {
                    "params": item.params,
                    "code": sub_bc.to_dict(),
                }
        self._emit("MAKE_CLASS", {
            "name": node.name,
            "base": node.base,
            "methods": methods,
        })
        self._emit("STORE_NAME", node.name)

    def _c_ImportStatement(self, node):
        self._emit("IMPORT_NAME", {"module": node.module, "alias": node.alias})

    def _e_TryExpr(self, node):
        # Simplified: emit TRY_BEGIN / TRY_END markers for VM
        try_start = self._emit("TRY_BEGIN", None)
        for stmt in node.try_block:
            self._compile_node(stmt)
        jump_past = self._emit("JUMP", None)
        catch_start = len(self.bc.instructions)
        self._patch(try_start, catch_start)
        if getattr(node, 'catch_bind', None):
            self._emit("STORE_NAME", node.catch_bind)
        else:
            self._emit("POP")
        for stmt in node.catch_block:
            self._compile_node(stmt)
        end = len(self.bc.instructions)
        self._patch(jump_past, end)
        if node.finally_block:
            for stmt in node.finally_block:
                self._compile_node(stmt)

    def _c_ThrowStatement(self, node):
        self._compile_expr(node.value)
        self._emit("RAISE")

    def _e_MatchExpr(self, node):
        self._compile_expr(node.subject)
        end_patches = []
        for case in node.cases:
            self._emit("DUP")  # keep subject on stack for each case
            self._compile_expr(case.pattern)
            eq_check = self._emit("JUMP_IF_NOT_EQ", None)
            self._emit("POP")  # pop subject when matched
            for stmt in case.body:
                self._compile_node(stmt)
            end_patches.append(self._emit("JUMP", None))
            self._patch(eq_check, len(self.bc.instructions))
        self._emit("POP")  # discard unmatched subject
        if node.default_body:
            for stmt in node.default_body:
                self._compile_node(stmt)
        for ep in end_patches:
            self._patch(ep, len(self.bc.instructions))

    def _e_BlockExpr(self, node):
        for stmt in node.stmts:
            self._compile_node(stmt)

    def _c_TestBlock(self, node):
        # Test blocks run unconditionally in bytecode mode
        self._emit("LOAD_CONST", f"\n  🧪 Running: \"{node.name}\"")
        self._emit("PRINT")
        for stmt in node.body:
            self._compile_node(stmt)

    def _c_InterfaceDef(self, node):
        self._emit("LOAD_CONST", {"__interface__": True, "methods": node.method_names})
        self._emit("STORE_NAME", node.name)

    def _c_ExportStatement(self, node):
        self._emit("LOAD_CONST", node.names)
        self._emit("STORE_NAME", "__exports__")

    def _c_AbstractMethodDef(self, node):
        self._emit("MAKE_FUNC", {"name": node.name, "params": node.params,
                                  "code": None, "is_abstract": True})
        self._emit("STORE_NAME", node.name)

    def _c_LoopStatement(self, node):
        self.__class__._c_LoopStatement(self, node)

    # ── Expression compilers ─────────────────────────────────────────────────
    def _e_NumberLiteral(self, node): self._emit("LOAD_CONST", node.value)
    def _e_StringLiteral(self, node): self._emit("LOAD_CONST", node.value)
    def _e_BoolLiteral(self, node):   self._emit("LOAD_CONST", node.value)
    def _e_NullLiteral(self, node):   self._emit("LOAD_CONST", None)
    def _e_SelfRef(self, node):       self._emit("LOAD_NAME", "self")
    def _e_SuperRef(self, node):      self._emit("LOAD_NAME", "__super__")

    def _e_Identifier(self, node):    self._emit("LOAD_NAME", node.name)

    def _e_BinaryOp(self, node):
        self._compile_expr(node.left)
        self._compile_expr(node.right)
        self._emit("BINARY_OP", node.op)

    def _e_UnaryOp(self, node):
        self._compile_expr(node.operand)
        self._emit("UNARY_OP", node.op)

    def _e_FunctionCall(self, node):
        self._compile_expr(node.callee)
        for arg in node.args:
            self._compile_expr(arg)
        for k, v in node.kwargs.items():
            self._compile_expr(v)
            self._emit("LOAD_CONST", k)
        self._emit("CALL", {"nargs": len(node.args), "nkwargs": len(node.kwargs)})

    def _e_MemberAccess(self, node):
        self._compile_expr(node.obj)
        self._emit("LOAD_ATTR", node.member)

    def _e_IndexAccess(self, node):
        self._compile_expr(node.obj)
        self._compile_expr(node.index)
        self._emit("LOAD_INDEX")

    def _e_ListLiteral(self, node):
        for el in node.elements:
            self._compile_expr(el)
        self._emit("BUILD_LIST", len(node.elements))

    def _e_DictLiteral(self, node):
        for k, v in node.pairs:
            self._compile_expr(k)
            self._compile_expr(v)
        self._emit("BUILD_DICT", len(node.pairs))

    def _e_RangeExpr(self, node):
        self._compile_expr(node.start)
        self._compile_expr(node.end)
        self._emit("BUILD_RANGE", node.inclusive)

    def _e_Assignment(self, node):
        self._compile_expr(node.value)
        self._emit("DUP")  # keep value on stack (assignment is expression)
        self._emit("STORE_NAME", node.name)

    def _e_NonNullAssertion(self, node):
        self._compile_expr(node.expr)
        self._emit("ASSERT_NOT_NULL")

    def _e_FString(self, node):
        parts = []
        for kind, val in node.parts:
            if kind == "lit":
                parts.append(("lit", val))
            else:
                self._compile_expr(val)
                parts.append(("expr", None))
        self._emit("BUILD_FSTRING", parts)

    def _e_LambdaExpr(self, node):
        sub = Compiler(self.bc.source)
        body_node = ReturnStatement(node.body) if not isinstance(node.body, list) else Program(node.body)
        sub_bc = sub.compile(Program([body_node]))
        self._emit("MAKE_FUNC", {"name": "<lambda>", "params": node.params,
                                  "code": sub_bc.to_dict(), "is_async": False})

    def _e_TernaryOp(self, node):
        self._compile_expr(node.condition)
        jump_false = self._emit("JUMP_IF_FALSE", None)
        self._compile_expr(node.true_expr)
        jump_end = self._emit("JUMP", None)
        self._patch(jump_false, len(self.bc.instructions))
        self._compile_expr(node.false_expr)
        self._patch(jump_end, len(self.bc.instructions))

    def _e_NullishCoalescing(self, node):
        self._compile_expr(node.left)
        self._emit("DUP")
        jump_if_not_null = self._emit("JUMP_IF_NOT_NULL", None)
        self._emit("POP")
        self._compile_expr(node.right)
        self._patch(jump_if_not_null, len(self.bc.instructions))
