"""
AVM ISA v1 — Python Port (Strict Functional Execution)
Mirrors formal/SilverSight/AVMIsa/Step.lean
"""
from __future__ import annotations
from dataclasses import dataclass, field
from enum import IntEnum, auto
from typing import Union, Optional
import sys

from q16_canonical import Q16_SCALE, float_to_q16, q16_to_float

# ── Constants ────────────────────────────────────────────────────────
AVM_CLAMP_MIN = -2147483647
AVM_CLAMP_MAX = 2147483647
AVM_Q0_MIN = -32767
AVM_Q0_MAX = 32767
AVM_MAX_STACK = 1024

def avm_clamp(x: int) -> int:
    return max(AVM_CLAMP_MIN, min(AVM_CLAMP_MAX, x))

def avm_q0_clamp(x: int) -> int:
    return max(AVM_Q0_MIN, min(AVM_Q0_MAX, x))

def floor_div(a: int, b: int) -> int:
    """Floor division matching Lean Int.ediv."""
    if b == 0:
        raise ValueError("division by zero")
    return -((-a) // b) if (a < 0) != (b < 0) and a % b != 0 else a // b

def lt_q16_v6(a: int, b: int) -> bool:
    sa = a < 0
    sb = b < 0
    return sa if sa != sb else a < b

# ── Types ────────────────────────────────────────────────────────────
class AvmTy(IntEnum):
    Q0_16 = 0
    Q16_16 = 1
    BOOL = 2

# ── Values ───────────────────────────────────────────────────────────
@dataclass
class AnyVal:
    ty: AvmTy
    val: Union[int, bool]

    @staticmethod
    def q16(x: int) -> 'AnyVal':
        return AnyVal(AvmTy.Q16_16, avm_clamp(x))
    @staticmethod
    def q0(x: int) -> 'AnyVal':
        return AnyVal(AvmTy.Q0_16, avm_q0_clamp(x))
    @staticmethod
    def b(x: bool) -> 'AnyVal':
        return AnyVal(AvmTy.BOOL, x)

# ── Instructions ─────────────────────────────────────────────────────
class Prim(IntEnum):
    ADD_SAT_Q0 = 0
    SUB_SAT_Q0 = 1
    ADD_SAT_Q16 = 2
    SUB_SAT_Q16 = 3
    MUL_SAT_Q16 = 4
    DIV_SAT_Q16 = 5
    LT_Q16 = 6
    EQ_Q16 = 7
    AND = 8
    OR = 9
    NOT = 10

    @property
    def arity(self) -> int:
        return 1 if self == Prim.NOT else 2

class Instr:
    PUSH_Q16, PUSH_BOOL, PUSH_Q0, POP, DUP, SWAP, LOAD, STORE, JUMP, JUMP_IF, PRIM, HALT = range(12)

# ── Primitive execution ──────────────────────────────────────────────
def eval_prim(p: Prim, a: AnyVal, b: Optional[AnyVal] = None) -> AnyVal:
    if p == Prim.ADD_SAT_Q0:
        assert a.ty == b.ty == AvmTy.Q0_16
        return AnyVal.q0(a.val + b.val)
    elif p == Prim.SUB_SAT_Q0:
        assert a.ty == b.ty == AvmTy.Q0_16
        return AnyVal.q0(a.val - b.val)
    elif p == Prim.ADD_SAT_Q16:
        assert a.ty == b.ty == AvmTy.Q16_16
        return AnyVal.q16(a.val + b.val)
    elif p == Prim.SUB_SAT_Q16:
        assert a.ty == b.ty == AvmTy.Q16_16
        return AnyVal.q16(a.val - b.val)
    elif p == Prim.MUL_SAT_Q16:
        assert a.ty == b.ty == AvmTy.Q16_16
        return AnyVal.q16(floor_div(a.val * b.val, Q16_SCALE))
    elif p == Prim.DIV_SAT_Q16:
        assert a.ty == b.ty == AvmTy.Q16_16
        if b.val == 0:
            raise ValueError("division by zero")
        return AnyVal.q16(floor_div(a.val * Q16_SCALE, b.val))
    elif p == Prim.LT_Q16:
        assert a.ty == b.ty == AvmTy.Q16_16
        return AnyVal.b(lt_q16_v6(a.val, b.val))
    elif p == Prim.EQ_Q16:
        assert a.ty == b.ty == AvmTy.Q16_16
        return AnyVal.b(a.val == b.val)
    elif p == Prim.AND:
        assert a.ty == b.ty == AvmTy.BOOL
        return AnyVal.b(a.val and b.val)
    elif p == Prim.OR:
        assert a.ty == b.ty == AvmTy.BOOL
        return AnyVal.b(a.val or b.val)
    elif p == Prim.NOT:
        assert a.ty == AvmTy.BOOL
        return AnyVal.b(not a.val)

# ── Errors ───────────────────────────────────────────────────────────
class StepError(Exception):
    def __init__(self, kind: str):
        self.kind = kind

# ── State ────────────────────────────────────────────────────────────
@dataclass
class State:
    pc: int = 0
    stack: list = field(default_factory=list)
    locals: list = field(default_factory=list)
    halted: bool = False

    @staticmethod
    def new(n_locals: int = 0):
        return State(pc=0, stack=[], locals=[None] * n_locals, halted=False)

# ── Step ─────────────────────────────────────────────────────────────
def step(s: State, prog: list) -> State:
    if s.halted:
        raise StepError("halted")
    if s.pc < 0 or s.pc >= len(prog):
        return State(pc=s.pc, stack=list(s.stack), locals=list(s.locals), halted=True)

    instr = prog[s.pc]
    stack = list(s.stack)
    pc = s.pc + 1
    halted = False

    GROW_OPS = {Instr.PUSH_Q16, Instr.PUSH_BOOL, Instr.PUSH_Q0, Instr.DUP, Instr.LOAD}
    if instr[0] in GROW_OPS and len(stack) >= AVM_MAX_STACK:
        raise StepError("stack_overflow")

    op = instr[0]
    arg = instr[1]
    arg2 = instr[2] if len(instr) > 2 else None

    if op == Instr.PUSH_Q16:
        stack.append(AnyVal.q16(arg))
    elif op == Instr.PUSH_BOOL:
        stack.append(AnyVal.b(arg2))
    elif op == Instr.PUSH_Q0:
        stack.append(AnyVal.q0(arg))
    elif op == Instr.POP:
        if not stack: raise StepError("empty_stack")
        stack.pop()
    elif op == Instr.DUP:
        if not stack: raise StepError("empty_stack")
        stack.append(stack[-1])
    elif op == Instr.SWAP:
        if len(stack) < 2: raise StepError("stack_underflow")
        stack[-1], stack[-2] = stack[-2], stack[-1]
    elif op == Instr.LOAD:
        if arg >= len(s.locals) or s.locals[arg] is None:
            raise StepError("missing_local")
        stack.append(s.locals[arg])
    elif op == Instr.STORE:
        if not stack: raise StepError("empty_stack")
        if arg >= len(s.locals): raise StepError("missing_local")
        s.locals[arg] = stack.pop()
    elif op == Instr.JUMP:
        if arg < 0 or arg >= len(prog): raise StepError("jump_out_of_bounds")
        pc = arg
    elif op == Instr.JUMP_IF:
        if not stack: raise StepError("empty_stack")
        cond = stack.pop()
        if cond.ty != AvmTy.BOOL: raise StepError("type_mismatch")
        if cond.val:
            if arg < 0 or arg >= len(prog): raise StepError("jump_out_of_bounds")
            pc = arg
    elif op == Instr.PRIM:
        p = Prim(arg)
        arity = p.arity
        if len(stack) < arity: raise StepError("stack_underflow")
        b = stack.pop() if arity >= 2 else None
        a = stack.pop()
        stack.append(eval_prim(p, a, b))
    elif op == Instr.HALT:
        halted = True
    else:
        raise StepError("unknown_instr")

    return State(pc=pc, stack=stack, locals=list(s.locals), halted=halted)

# ── Run (fuel-bounded) ──────────────────────────────────────────────
def run(initial: State, prog: list, fuel: int = 10000) -> State:
    s = initial
    for _ in range(fuel):
        if s.halted:
            return s
        s = step(s, prog)
    return s