"""The HDC virtual machine.

State is hyperdimensional and lives in storage:

* 16 registers, each one hypervector;
* the program counter indexes cells ``code/<pc>`` in HDC storage;
* the stack holds hypervectors;
* the flag is derived from a comparison, not stored in a status word;
* integer values are themselves hypervectors (the bit-position codec in
  :mod:`hdc.hypervec`), so ``LDI R0, 5`` writes a *vector*, and ``ADD``
  decodes two vectors to integers, adds them with the binary ALU, and
  re-encodes the sum.  That is the honest split in this design: HDC carries
  addressing, dispatch, records and associative recall natively; scalar
  arithmetic borrows a conventional integer ALU on the decoded payload.
  ``BIND``/``BUNDLE``/``PERMUTE``/``SIM``/``CLEAN``/``RECALL`` never leave
  hyperdimensional space at all.

Every fetched instruction is decoded from its vector alone — no side table of
"what the instruction was" is consulted, only cleanup alphabets supplied by
the assembled image.  Decodes are memoised by vector value, which is sound
because decoding is a pure function of the vector and the alphabet.
"""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable, Iterable

from .hypervec import (
    DEFAULT_DIM,
    bind,
    bits,
    bundle,
    hvc,
    hv_to_int,
    int_to_hv,
    permute,
    popcount,
    sim,
)
from .adder import HyperAdder
from .isa import (
    OPCODES,
    REG_COUNT,
    REG_NAMES,
    DecodeContext,
    DecodedInstr,
    decode_instr,
    format_operands,
)
from .memory import HDCStorage

#: Sentinel for "ran past the end of the program".
PROGRAM_END = -1


class MachineError(Exception):
    """A runtime fault inside the HDC machine."""


@dataclass
class TraceEntry:
    """One executed instruction, as the debugger wants to see it."""

    pc: int
    op: str
    operands: str
    confidence: float
    note: str = ""
    written: tuple[str, ...] = ()
    fetch_sim: float = 1.0

    def render(self) -> str:
        tail = ("  ; " + self.note) if self.note else ""
        return "%04d  %-8s %-26s conf=%.2f%s" % (
            self.pc,
            self.op,
            self.operands,
            self.confidence,
            tail,
        )


class Machine:
    """Fetch from HDC storage, decode by cleanup, execute, keep everything HDC."""

    def __init__(
        self,
        storage: HDCStorage | None = None,
        d: int = DEFAULT_DIM,
        tools=None,
        echo: Callable[[str], None] | None = None,
        adder_radix: int = 8,
    ):
        self.d = d
        self.storage = storage if storage is not None else HDCStorage(d=d)
        self.tools = tools
        self.echo = echo
        self.prefix = "code"
        self.prog_len = 0
        self.ctx = DecodeContext(d=d)
        #: ADD and SUB are executed by this circuit, not by the integer ALU.
        #: Its truth table lives in HDC storage of its own, so a program's own
        #: cell usage is unaffected.  radix=8 means 4 byte-slice stages per
        #: 32-bit word instead of 32 single-bit stages.
        self.adder = HyperAdder(d=d, radix=adder_radix)
        self.regs: list[int] = [0] * REG_COUNT
        self.pc = 0
        self.halted = False
        self.status = "ready"
        self.flag_z = False
        self.stack: list[int] = []
        self.output: list[str] = []
        self.steps = 0
        self.last_recall: list[tuple[str, int, float]] = []
        self.last_tool_result: object = None
        self.last_fetch_sim = 1.0
        self.skip_break = False
        self.trace: list[TraceEntry] = []
        self.trace_limit = 512
        self._decode_cache: dict[int, DecodedInstr] = {}
        self.reset()

    # ------------------------------------------------------------- lifecycle
    def reset(self) -> None:
        """Clear registers, stack, flags and counters (not storage)."""
        self.regs = [int_to_hv(0, self.d) for _ in range(REG_COUNT)]
        self.pc = 0
        self.halted = False
        self.status = "ready"
        self.flag_z = False
        self.stack = []
        self.output = []
        self.steps = 0
        self.trace = []
        self.last_recall = []
        self.last_tool_result = None

    def load(self, program, prefix: str = "code") -> "Machine":
        """Install an assembled program into storage and point the PC at 0."""
        self.prefix = prefix
        program.install(self.storage, prefix)
        self.prog_len = len(program.instrs)
        self.ctx = program.context()
        self._decode_cache.clear()
        self.reset()
        self.status = "ready"
        return self

    # ----------------------------------------------------------- fetch path
    def cell_name(self, pc: int) -> str:
        return "%s/%d" % (self.prefix, pc)

    def fetch(self, pc: int | None = None) -> int:
        """Read the instruction vector for a program cell out of HDC storage.

        Resolution goes through the content-addressable path: an exact address
        hit is normal, but a modified address vector still resolves and the
        similarity is reported to the debugger.
        """
        if pc is None:
            pc = self.pc
        if pc < 0 or (self.prog_len and pc >= self.prog_len):
            raise MachineError("PC out of program: %d" % pc)
        name = self.cell_name(pc)
        resolved = self.storage.resolve(self.storage.addr(name))
        if resolved is None:
            raise MachineError("no instruction cell at pc=%d" % pc)
        cell, confidence = resolved
        self.last_fetch_sim = confidence
        return self.storage.read(cell)

    def decode_at(self, pc: int | None = None) -> DecodedInstr:
        """Fetch and decode one instruction."""
        hv = self.fetch(pc)
        dec = self._decode_cache.get(hv)
        if dec is None:
            dec = decode_instr(hv, self.ctx)
            self._decode_cache[hv] = dec
        return dec

    # --------------------------------------------------------------- values
    def get_int(self, idx: int) -> int:
        """Decode a register's integer payload."""
        return hv_to_int(self.regs[idx], self.d)

    def set_int(self, idx: int, value: int) -> None:
        """Encode an integer back into a register as a hypervector."""
        self.regs[idx] = int_to_hv(value, self.d)

    def get_sim(self, a: int, b: int) -> float:
        return sim(self.regs[a], self.regs[b], self.d)

    # -------------------------------------------------------------- execute
    def step(self) -> TraceEntry | None:
        """Execute exactly one instruction; returns its trace entry."""
        if self.halted:
            return None
        start_pc = self.pc
        dec = self.decode_at(start_pc)
        op = dec.op
        note = ""
        written: list[str] = []
        next_pc = start_pc + 1

        def dst(i: int, value: int, tag: str = "") -> None:
            if 0 <= i < REG_COUNT:
                self.regs[i] = value
                written.append(REG_NAMES[i])

        try:
            if op == "NOP":
                pass
            elif op == "HALT":
                self.halted = True
                self.status = "halted"
                note = "halt"
            elif op == "LDI":
                dst(dec.regs[0], int_to_hv(int(dec.imm or 0), self.d))
                note = "imm=%d" % (dec.imm or 0)
            elif op == "MOV":
                dst(dec.regs[0], self.regs[dec.regs[1]])
            elif op in ("ADD", "SUB"):
                # Native hyperdimensional arithmetic: 32 stages of bind + cleanup
                # against the stored full-adder table.  No integer payload is
                # materialised to compute this result; the operands stay vectors
                # from fetch to store.
                a, b, c = dec.regs[0], dec.regs[1], dec.regs[2]
                if op == "ADD":
                    dst(a, self.adder.add_hv(self.regs[b], self.regs[c]))
                else:
                    dst(a, self.adder.sub_hv(self.regs[b], self.regs[c]))
                note = "%s: %d stages, radix %d" % (
                    op.lower(), len(self.adder.last_stages), self.adder.radix
                )
            elif op in ("MUL", "DIV", "AND", "OR", "XOR"):
                a, b, c = dec.regs[0], dec.regs[1], dec.regs[2]
                x, y = self.get_int(b), self.get_int(c)
                if op == "MUL":
                    r = x * y
                elif op == "DIV":
                    r = x // y if y else 0
                elif op == "AND":
                    r = x & y
                elif op == "OR":
                    r = x | y
                else:
                    r = x ^ y
                dst(a, int_to_hv(r, self.d))
                note = "%d %s %d = %d" % (x, op.lower(), y, r)
            elif op == "NOT":
                x = self.get_int(dec.regs[1])
                dst(dec.regs[0], int_to_hv(~x, self.d))
                note = "~%d" % x
            elif op in ("SHL", "SHR"):
                x = self.get_int(dec.regs[1])
                n = int(dec.imm or 0)
                r = x << n if op == "SHL" else x >> n
                dst(dec.regs[0], int_to_hv(r, self.d))
                note = "%d %s %d = %d" % (x, "<<" if op == "SHL" else ">>", n, r)
            elif op == "POPC":
                dst(dec.regs[0], int_to_hv(popcount(self.regs[dec.regs[1]]), self.d))
                note = "popcount"
            elif op == "CMP":
                self.flag_z = self.get_int(dec.regs[0]) == self.get_int(dec.regs[1])
                note = "ZF=%d" % int(self.flag_z)
            elif op == "JMP":
                next_pc = int(dec.tgt or 0)
            elif op == "JZ":
                next_pc = int(dec.tgt or 0) if self.flag_z else next_pc
            elif op == "JNZ":
                next_pc = int(dec.tgt or 0) if not self.flag_z else next_pc
            elif op == "CALL":
                self.stack.append(int_to_hv(next_pc, self.d))
                next_pc = int(dec.tgt or 0)
                note = "call"
            elif op == "RET":
                if not self.stack:
                    raise MachineError("RET with empty stack")
                next_pc = hv_to_int(self.stack.pop(), self.d)
                note = "ret"
            elif op == "PUSH":
                self.stack.append(self.regs[dec.regs[0]])
            elif op == "POP":
                if not self.stack:
                    raise MachineError("POP with empty stack")
                dst(dec.regs[0], self.stack.pop())
            elif op == "BIND":
                dst(dec.regs[0], bind(self.regs[dec.regs[1]], self.regs[dec.regs[2]], self.d))
                note = "bind"
            elif op == "BUNDLE":
                dst(
                    dec.regs[0],
                    bundle([self.regs[dec.regs[1]], self.regs[dec.regs[2]]], self.d),
                )
                note = "bundle"
            elif op == "PERMUTE":
                dst(dec.regs[0], permute(self.regs[dec.regs[1]], int(dec.imm or 0), self.d))
                note = "rho^%d" % int(dec.imm or 0)
            elif op == "SIM":
                s = self.get_sim(dec.regs[1], dec.regs[2])
                dst(dec.regs[0], int_to_hv(int(round(s * 1000)), self.d))
                note = "sim=%.4f" % s
            elif op == "CLEAN":
                hits = self.storage.nearest_symbol(self.regs[dec.regs[1]], k=1)
                if hits:
                    sym_name, sym_sim = hits[0]
                    symbols = self.storage.symbols()
                    idx = symbols.index(sym_name) if sym_name in symbols else -1
                    note = "%s sim=%.3f" % (sym_name, sym_sim)
                else:
                    idx, note = -1, "no symbol"
                dst(dec.regs[0], int_to_hv(idx, self.d))
            elif op == "STORE":
                if dec.name is None:
                    raise MachineError("STORE needs a cell name")
                self.storage.write(dec.name, self.regs[dec.regs[0]], symbol=True)
                note = "->%s" % dec.name
            elif op == "LOAD":
                if dec.name is None:
                    raise MachineError("LOAD needs a cell name")
                dst(dec.regs[0], self.storage.read(dec.name, default=int_to_hv(0, self.d)))
                note = "<-%s" % dec.name
            elif op == "RECALL":
                k = int(dec.imm or 1)
                probe = self.regs[dec.regs[1]]
                hits = self.storage.recall(probe, k=max(1, k))
                self.last_recall = hits
                best = hits[0] if hits else None
                order = self.storage.order
                idx = order.index(best[0]) if best and best[0] in order else -1
                dst(dec.regs[0], int_to_hv(idx, self.d))
                note = "%s sim=%.3f" % (best[0], best[2]) if best else "no hit"
            elif op == "MEMSTAT":
                dst(dec.regs[0], int_to_hv(self.storage.stat()["cells_used"], self.d))
            elif op == "OUT":
                v = self.get_int(dec.regs[0])
                self._emit(str(v))
                note = "out %d" % v
            elif op == "OUTT":
                hits = self.storage.nearest_symbol(self.regs[dec.regs[0]], k=1)
                if hits:
                    self._emit("%s  (sim %.3f)" % (hits[0][0], hits[0][1]))
                    note = "outt %s" % hits[0][0]
                else:
                    self._emit("<no symbol>")
            elif op == "PRINTS":
                self._emit(str(dec.text or ""))
                note = "prints"
            elif op == "TOOL":
                self._tool(dec, dst)
                note = "tool=%s" % dec.tool
            elif op == "BREAK":
                if self.skip_break:
                    self.skip_break = False
                    note = "breakpoint stepped over"
                else:
                    self.status = "break"
                    note = "breakpoint trap"
                    self.trace.append(
                        TraceEntry(start_pc, op, format_operands(dec), dec.confidence, note,
                                   tuple(written), self.last_fetch_sim)
                    )
                    return self.trace[-1]
            else:  # pragma: no cover - unreachable while OPCODES is closed
                raise MachineError("unimplemented opcode %r" % op)
        except MachineError as exc:
            self.status = "error"
            self.halted = True
            note = "FAULT: %s" % exc

        self.steps += 1
        self.pc = next_pc
        if not self.halted and self.prog_len and self.pc >= self.prog_len:
            self.halted = True
            self.status = "halted"
        entry = TraceEntry(
            start_pc, op, format_operands(dec), dec.confidence, note, tuple(written),
            self.last_fetch_sim,
        )
        self.trace.append(entry)
        if len(self.trace) > self.trace_limit:
            del self.trace[: len(self.trace) - self.trace_limit]
        return entry

    def _emit(self, text: str) -> None:
        self.output.append(text)
        if self.echo is not None:
            self.echo(text)

    def _tool(self, dec: DecodedInstr, dst) -> None:
        if self.tools is None:
            raise MachineError("TOOL executed but no tool registry is attached")
        name = dec.tool or ""
        try:
            params = self.tools.params_for(name)
            ret_kind = self.tools.returns_for(name)
            actual = [dec.regs[1], dec.regs[2]][: len(params)]
            args = []
            for kind, reg in zip(params, actual):
                args.append(self.get_int(reg) if kind == "int" else self.regs[reg])
            result = self.tools.call(name, args, machine=self)
        except Exception as exc:
            # A failed tool is a machine fault the program can be debugged
            # through, not an exception escaping the VM.
            raise MachineError("tool %r failed: %s" % (name, exc)) from exc
        self.last_tool_result = result
        if isinstance(result, bool):
            result = int(result)
        if ret_kind == "hv":
            dst(dec.regs[0], int(result) & ((1 << self.d) - 1))
        elif ret_kind == "str":
            dst(dec.regs[0], hvc("str/" + str(result), self.d))
        elif isinstance(result, int):
            dst(dec.regs[0], int_to_hv(int(result), self.d))
        else:
            raise MachineError("tool %r returned an unusable type" % name)

    # ------------------------------------------------------------------ run
    def run(
        self,
        max_steps: int = 100_000,
        on_step: Callable[["Machine", TraceEntry], None] | None = None,
    ) -> dict:
        """Run until HALT, a breakpoint trap, a fault, or the step budget."""
        # Resuming after a trap or an exhausted budget must clear the parked
        # status, otherwise the first step would immediately stop again.
        if self.status in ("ready", "break", "budget"):
            self.status = "running"
        # Note: resuming past a trap cell is the caller's call — the debugger
        # and the tests set ``skip_break`` explicitly, so a bare re-run of a
        # trapped machine traps again rather than silently skipping.
        while not self.halted and self.steps < max_steps:
            entry = self.step()
            if entry is None:
                break
            if on_step is not None:
                on_step(self, entry)
            # Stop only on a real trap.  A BREAK that was deliberately
            # stepped over (skip_break) must not end the run.
            if entry.op == "BREAK" and self.status == "break":
                break
        if self.steps >= max_steps and not self.halted:
            self.status = "budget"
        return self.result()

    def result(self) -> dict:
        return {
            "status": self.status,
            "pc": self.pc,
            "steps": self.steps,
            "halted": self.halted,
            "flag_z": self.flag_z,
            "output": list(self.output),
            "stack_depth": len(self.stack),
        }

    # ---------------------------------------------------------- inspection
    def regs_text(self, with_values: bool = True) -> str:
        """Register dump: index, integer payload, population count."""
        rows = ["REG      VALUE       POP   HV(head)"]
        for i, hv in enumerate(self.regs):
            val = self.get_int(i) if with_values else 0
            rows.append(
                "%-8s %-11d %-5d %s"
                % (REG_NAMES[i], val, popcount(hv), bits(hv, self.d)[:32])
            )
        return "\n".join(rows)

    def disasm(self, start: int = 0, count: int = 12) -> str:
        """Disassemble program cells straight out of storage."""
        rows = []
        for pc in range(start, min(start + count, self.prog_len)):
            dec = self.decode_at(pc)
            mark = "=>" if pc == self.pc else "  "
            rows.append(
                "%s %04d  %-8s %-26s conf=%.2f"
                % (mark, pc, dec.op, format_operands(dec), dec.confidence)
            )
        return "\n".join(rows) if rows else "<empty program>"