"""
AVM ISA v1 — Python Port
Strict functional: step(state, program) -> state, run(state, program, fuel) -> state
"""
from dataclasses import dataclass, field
from typing import List, Union, Tuple

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

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

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

def floor_div(a, b):
    if b == 0: raise ValueError("div by zero")
    q, r = divmod(a, b)
    return q - 1 if r and (a ^ b) < 0 else q

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

def q16_mul(a, b):
    return floor_div(a * b, Q16_SCALE)

def q16_div(a, b):
    if b == 0: raise ValueError("div by zero")
    return floor_div(a * Q16_SCALE, b)

# ── Types ────────────────────────────────────────────────────────────
class Prim:
    ADD_SAT_Q0, SUB_SAT_Q0, ADD_SAT_Q16, SUB_SAT_Q16, \
        MUL_SAT_Q16, DIV_SAT_Q16, LT_Q16, EQ_Q16, AND, OR, NOT = range(11)

prim_arity = {Prim.NOT: 1}.get  # all others are 2

class AvmVal:
    def __init__(self, ty, val):
        self.ty = ty
        self.val = val
    @staticmethod
    def q0(x): return AvmVal('q0', avm_q0_clamp(x))
    @staticmethod
    def q16(x): return AvmVal('q16', avm_clamp(x))
    @staticmethod
    def bool(x): return AvmVal('bool', bool(x))

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

Q0_OPS = {Prim.ADD_SAT_Q0, Prim.SUB_SAT_Q0}
Q16_BIN_OPS = {Prim.ADD_SAT_Q16, Prim.SUB_SAT_Q16, Prim.MUL_SAT_Q16, Prim.DIV_SAT_Q16, Prim.LT_Q16, Prim.EQ_Q16}
BOOL_BIN_OPS = {Prim.AND, Prim.OR}

def check_type(val_ty, expected, label):
    if val_ty != expected:
        raise TypeError(f"type mismatch: expected {expected}, got {val_ty}")

def make_instr(op, arg=None, arg2=None):
    return {'op': op, 'arg': arg, 'arg2': arg2}

def push_q16(x): return make_instr(Instr.PUSH, avm_clamp(x))
def push_bool(x): return make_instr(Instr.PUSH_BOOL, 0, x)
def push_q0(x): return make_instr(Instr.PUSH_Q0, avm_q0_clamp(x))
def pop(): return make_instr(Instr.POP)
def dup(): return make_instr(Instr.DUP)
def swap(): return make_instr(Instr.SWAP)
def load(i): return make_instr(Instr.LOAD, i)
def store(i): return make_instr(Instr.STORE, i)
def jump(t): return make_instr(Instr.JUMP, t)
def jump_if(t): return make_instr(Instr.JUMP_IF, t)
def prim(p): return make_instr(Instr.PRIM, p)
def halt(): return make_instr(Instr.HALT)

@dataclass
class State:
    pc: int = 0
    stack: list = field(default_factory=list)
    locals: list = field(default_factory=lambda: [None] * 16)
    halted: bool = False

# ── Primitive execution ──────────────────────────────────────────────
def exec_prim(op, a, b=None):
    if op == Prim.ADD_SAT_Q0:
        check_type(a.ty, 'q0', 'a'); check_type(b.ty, 'q0', 'b')
        return AvmVal.q0(a.val + b.val)
    if op == Prim.SUB_SAT_Q0:
        check_type(a.ty, 'q0', 'a'); check_type(b.ty, 'q0', 'b')
        return AvmVal.q0(a.val - b.val)
    if op == Prim.ADD_SAT_Q16:
        check_type(a.ty, 'q16', 'a'); check_type(b.ty, 'q16', 'b')
        return AvmVal.q16(a.val + b.val)
    if op == Prim.SUB_SAT_Q16:
        check_type(a.ty, 'q16', 'a'); check_type(b.ty, 'q16', 'b')
        return AvmVal.q16(a.val - b.val)
    if op == Prim.MUL_SAT_Q16:
        check_type(a.ty, 'q16', 'a'); check_type(b.ty, 'q16', 'b')
        return AvmVal.q16(q16_mul(a.val, b.val))
    if op == Prim.DIV_SAT_Q16:
        check_type(a.ty, 'q16', 'a'); check_type(b.ty, 'q16', 'b')
        if b.val == 0: raise ValueError("division by zero")
        return AvmVal.q16(q16_div(a.val, b.val))
    if op == Prim.LT_Q16:
        check_type(a.ty, 'q16', 'a'); check_type(b.ty, 'q16', 'b')
        return AvmVal.bool(lt_q16_v6(a.val, b.val))
    if op == Prim.EQ_Q16:
        check_type(a.ty, 'q16', 'a'); check_type(b.ty, 'q16', 'b')
        return AvmVal.bool(a.val == b.val)
    if op == Prim.AND:
        check_type(a.ty, 'bool', 'a'); check_type(b.ty, 'bool', 'b')
        return AvmVal.bool(a.val and b.val)
    if op == Prim.OR:
        check_type(a.ty, 'bool', 'a'); check_type(b.ty, 'bool', 'b')
        return AvmVal.bool(a.val or b.val)
    if op == Prim.NOT:
        check_type(a.ty, 'bool', 'a')
        return AvmVal.bool(not a.val)
    raise TypeError(f"unknown prim {op}")

# ── Step ─────────────────────────────────────────────────────────────
def step(state, program):
    if state.halted: raise ValueError("halted")
    if state.pc < 0 or state.pc >= len(program):
        return State(state.pc, list(state.stack), list(state.locals), True)

    instr = program[state.pc]
    stack = list(state.stack)
    locals = list(state.locals)
    pc = state.pc + 1
    halted = False

    growing = instr['op'] in (Instr.PUSH, Instr.PUSH_BOOL, Instr.PUSH_Q0, Instr.DUP, Instr.LOAD)
    if growing and len(stack) >= AVM_MAX_STACK:
        raise ValueError("stack overflow")

    if instr['op'] == Instr.PUSH:
        stack.append(AvmVal.q16(instr['arg']))
    elif instr['op'] == Instr.PUSH_BOOL:
        stack.append(AvmVal.bool(instr['arg2']))
    elif instr['op'] == Instr.PUSH_Q0:
        stack.append(AvmVal.q0(instr['arg']))
    elif instr['op'] == Instr.POP:
        if not stack: raise ValueError("empty stack")
        stack.pop()
    elif instr['op'] == Instr.DUP:
        if not stack: raise ValueError("empty stack")
        stack.append(stack[-1])
    elif instr['op'] == Instr.SWAP:
        if len(stack) < 2: raise ValueError("stack underflow")
        stack[-1], stack[-2] = stack[-2], stack[-1]
    elif instr['op'] == Instr.LOAD:
        i = instr['arg']
        if i >= len(locals) or locals[i] is None: raise ValueError("missing local")
        stack.append(locals[i])
    elif instr['op'] == Instr.STORE:
        if not stack: raise ValueError("empty stack")
        i = instr['arg']
        if i >= len(locals): raise ValueError("missing local")
        locals[i] = stack.pop()
    elif instr['op'] == Instr.JUMP:
        if instr['arg'] < 0 or instr['arg'] >= len(program): raise ValueError("jump OOB")
        pc = instr['arg']
    elif instr['op'] == Instr.JUMP_IF:
        if not stack: raise ValueError("empty stack")
        v = stack.pop()
        if v.ty != 'bool': raise TypeError("type mismatch")
        if v.val:
            if instr['arg'] < 0 or instr['arg'] >= len(program): raise ValueError("jump OOB")
            pc = instr['arg']
    elif instr['op'] == Instr.PRIM:
        p = instr['arg']
        arity = 1 if p == Prim.NOT else 2
        if len(stack) < arity: raise ValueError("stack underflow")
        b = stack.pop() if arity >= 2 else None
        a = stack.pop()
        stack.append(exec_prim(p, a, b))
    elif instr['op'] == Instr.HALT:
        halted = True
    else:
        raise ValueError(f"unknown instr {instr['op']}")

    return State(pc, stack, locals, halted)

# ── Run (fuel-bounded) ──────────────────────────────────────────────
def run(initial, program, fuel=1000):
    state = initial
    for _ in range(fuel):
        if state.halted: return state
        state = step(state, program)
    return state
