// AVM ISA v1 — C++ Port (Strict Functional Execution)
#pragma once
#include <cstdint>
#include <vector>
#include <variant>
#include <optional>
#include <stdexcept>

namespace avm {

// ── Constants ────────────────────────────────────────────────────
constexpr int32_t AVM_CLAMP_MIN = -2147483647;
constexpr int32_t AVM_CLAMP_MAX = 2147483647;
constexpr int32_t AVM_Q0_MIN = -32767;
constexpr int32_t AVM_Q0_MAX = 32767;
constexpr int64_t Q16_SCALE = 65536;
constexpr size_t AVM_MAX_STACK = 1024;

inline int32_t avm_clamp(int64_t x) {
    if (x > AVM_CLAMP_MAX) return AVM_CLAMP_MAX;
    if (x < AVM_CLAMP_MIN) return AVM_CLAMP_MIN;
    return static_cast<int32_t>(x);
}

inline int32_t avm_q0_clamp(int64_t x) {
    if (x > AVM_Q0_MAX) return AVM_Q0_MAX;
    if (x < AVM_Q0_MIN) return AVM_Q0_MIN;
    return static_cast<int32_t>(x);
}

inline int32_t floor_div(int64_t a, int64_t b) {
    if (b == 0) throw std::runtime_error("division by zero");
    int64_t q = a / b;
    if (a % b != 0 && ((a ^ b) < 0)) q--;
    return static_cast<int32_t>(q);
}

inline bool lt_q16_v6(int32_t a, int32_t b) {
    bool sa = a < 0, sb = b < 0;
    return (sa != sb) ? sa : (a < b);
}

// ── Types ──────────────────────────────────────────────────────
enum class Ty : uint8_t { Q0_16, Q16_16, Bool };
using Val = std::variant<int32_t, bool>;

struct AnyVal { Ty ty; Val val; };

// ── Primitives ────────────────────────────────────────────────
enum class Prim : uint8_t {
    AddSatQ0, SubSatQ0,
    AddSatQ16, SubSatQ16, MulSatQ16, DivSatQ16,
    LtQ16, EqQ16, And, Or, Not
};

inline int prim_arity(Prim p) {
    return (p == Prim::Not) ? 1 : 2;
}

// ── Instructions ──────────────────────────────────────────────
enum class Op : uint8_t {
    PushQ16, PushBool, PushQ0,
    Pop, Dup, Swap, Load, Store,
    Jump, JumpIf, Primitive, Halt
};

struct Instr {
    Op op;
    int32_t arg;
    bool arg2;
};

// ── Primitive execution ──────────────────────────────────────
inline AnyVal exec_prim(Prim p, const AnyVal& a, const AnyVal& b) {
    auto check = [](const AnyVal& v, Ty t) { if (v.ty != t) throw std::runtime_error("type mismatch"); };
    switch (p) {
        case Prim::AddSatQ0:
            check(a, Ty::Q0_16); check(b, Ty::Q0_16);
            return {Ty::Q0_16, avm_q0_clamp(static_cast<int64_t>(std::get<int32_t>(a.val)) + std::get<int32_t>(b.val))};
        case Prim::SubSatQ0:
            check(a, Ty::Q0_16); check(b, Ty::Q0_16);
            return {Ty::Q0_16, avm_q0_clamp(static_cast<int64_t>(std::get<int32_t>(a.val)) - std::get<int32_t>(b.val))};
        case Prim::AddSatQ16:
            check(a, Ty::Q16_16); check(b, Ty::Q16_16);
            return {Ty::Q16_16, avm_clamp(static_cast<int64_t>(std::get<int32_t>(a.val)) + std::get<int32_t>(b.val))};
        case Prim::SubSatQ16:
            check(a, Ty::Q16_16); check(b, Ty::Q16_16);
            return {Ty::Q16_16, avm_clamp(static_cast<int64_t>(std::get<int32_t>(a.val)) - std::get<int32_t>(b.val))};
        case Prim::MulSatQ16:
            check(a, Ty::Q16_16); check(b, Ty::Q16_16);
            return {Ty::Q16_16, avm_clamp(floor_div(static_cast<int64_t>(std::get<int32_t>(a.val)) * std::get<int32_t>(b.val), Q16_SCALE))};
        case Prim::DivSatQ16:
            check(a, Ty::Q16_16); check(b, Ty::Q16_16);
            return {Ty::Q16_16, avm_clamp(floor_div(static_cast<int64_t>(std::get<int32_t>(a.val)) * Q16_SCALE, std::get<int32_t>(b.val)))};
        case Prim::LtQ16:
            check(a, Ty::Q16_16); check(b, Ty::Q16_16);
            return {Ty::Bool, lt_q16_v6(std::get<int32_t>(a.val), std::get<int32_t>(b.val))};
        case Prim::EqQ16:
            check(a, Ty::Q16_16); check(b, Ty::Q16_16);
            return {Ty::Bool, std::get<int32_t>(a.val) == std::get<int32_t>(b.val)};
        case Prim::And:
            check(a, Ty::Bool); check(b, Ty::Bool);
            return {Ty::Bool, std::get<bool>(a.val) && std::get<bool>(b.val)};
        case Prim::Or:
            check(a, Ty::Bool); check(b, Ty::Bool);
            return {Ty::Bool, std::get<bool>(a.val) || std::get<bool>(b.val)};
        case Prim::Not:
            check(a, Ty::Bool);
            return {Ty::Bool, !std::get<bool>(a.val)};
    }
    throw std::runtime_error("unknown prim");
}

// ── State ─────────────────────────────────────────────────────
struct State {
    int pc = 0;
    std::vector<AnyVal> stack;
    std::vector<std::optional<AnyVal>> locals;
    bool halted = false;
};

inline State init_state(size_t n_locals = 0) {
    return {0, {}, std::vector<std::optional<AnyVal>>(n_locals), false};
}

// ── Step ─────────────────────────────────────────────────────
inline std::optional<State> step(const State& s, const std::vector<Instr>& prog) {
    if (s.halted) return std::nullopt;
    if (s.pc < 0 || static_cast<size_t>(s.pc) >= prog.size())
        return State{s.pc, s.stack, s.locals, true};

    auto instr = prog[s.pc];
    State ns = s;
    int npc = s.pc + 1;

    auto growing = (instr.op == Op::PushQ16 || instr.op == Op::PushBool || instr.op == Op::PushQ0 || instr.op == Op::Dup || instr.op == Op::Load);
    if (growing && ns.stack.size() >= AVM_MAX_STACK) return std::nullopt; // overflow

    switch (instr.op) {
        case Op::PushQ16: ns.stack.push_back({Ty::Q16_16, avm_clamp(instr.arg)}); break;
        case Op::PushBool: ns.stack.push_back({Ty::Bool, instr.arg2}); break;
        case Op::PushQ0: ns.stack.push_back({Ty::Q0_16, avm_q0_clamp(instr.arg)}); break;
        case Op::Pop:
            if (ns.stack.empty()) return std::nullopt;
            ns.stack.pop_back(); break;
        case Op::Dup:
            if (ns.stack.empty()) return std::nullopt;
            ns.stack.push_back(ns.stack.back()); break;
        case Op::Swap:
            if (ns.stack.size() < 2) return std::nullopt;
            std::swap(ns.stack[ns.stack.size()-1], ns.stack[ns.stack.size()-2]); break;
        case Op::Load: {
            size_t i = instr.arg;
            if (i >= ns.locals.size() || !ns.locals[i].has_value()) return std::nullopt;
            ns.stack.push_back(ns.locals[i].value()); break;
        }
        case Op::Store: {
            size_t i = instr.arg;
            if (ns.stack.empty() || i >= ns.locals.size()) return std::nullopt;
            ns.locals[i] = ns.stack.back(); ns.stack.pop_back(); break;
        }
        case Op::Jump:
            if (instr.arg < 0 || static_cast<size_t>(instr.arg) >= prog.size()) return std::nullopt;
            npc = instr.arg; break;
        case Op::JumpIf: {
            if (ns.stack.empty()) return std::nullopt;
            auto v = ns.stack.back(); ns.stack.pop_back();
            if (v.ty != Ty::Bool) return std::nullopt;
            if (std::get<bool>(v.val)) {
                if (instr.arg < 0 || static_cast<size_t>(instr.arg) >= prog.size()) return std::nullopt;
                npc = instr.arg;
            }
            break;
        }
        case Op::Primitive: {
            auto p = static_cast<Prim>(instr.arg);
            int arity = prim_arity(p);
            if (static_cast<int>(ns.stack.size()) < arity) return std::nullopt;
            AnyVal b{Ty::Bool, false};
            if (arity >= 2) { b = ns.stack.back(); ns.stack.pop_back(); }
            AnyVal a = ns.stack.back(); ns.stack.pop_back();
            try { ns.stack.push_back(exec_prim(p, a, b)); }
            catch (...) { return std::nullopt; }
            break;
        }
        case Op::Halt: ns.halted = true; break;
    }
    ns.pc = npc;
    return ns;
}

// ── Run (fuel-bounded) ───────────────────────────────────────
inline std::optional<State> run(const State& init, const std::vector<Instr>& prog, int fuel = 10000) {
    State s = init;
    for (int i = 0; i < fuel; i++) {
        if (s.halted) return s;
        auto next = step(s, prog);
        if (!next.has_value()) return std::nullopt;
        s = next.value();
    }
    return s;
}

} // namespace avm