// AVM ISA v1 — Scala Port (Strict Functional Execution)
package avm

object AVM {
  val Q16Scale: Int = 65536

  def q16Mul(a: Int, b: Int): Int = ((a.toLong * b.toLong) / Q16Scale).toInt
  def q16Div(a: Int, b: Int): Int = {
    if (b == 0) throw new ArithmeticException("div by zero")
    ((a.toLong * Q16Scale) / b).toInt
  }

  // ── Types ───────────────────────────────────────
  sealed trait ValType
  case object Q0 extends ValType
  case object Q16 extends ValType
  case object BoolT extends ValType

  case class AvmVal(ty: ValType, raw: Int)
  object AvmVal {
    def q16(x: Int): AvmVal = AvmVal(Q16, x)
    def bool(x: Boolean): AvmVal = AvmVal(BoolT, if (x) 1 else 0)
  }

  // ── Primitives ──────────────────────────────────
  sealed trait Prim
  case object AddQ16 extends Prim; case object SubQ16 extends Prim
  case object MulQ16 extends Prim; case object DivQ16 extends Prim
  case object LtQ16 extends Prim; case object EqQ16 extends Prim
  case object And extends Prim; case object Or extends Prim; case object Not extends Prim

  def execPrim(p: Prim, a: AvmVal, b: Option[AvmVal]): AvmVal = (p, b) match {
    case (AddQ16, Some(bv)) => AvmVal.q16(a.raw + bv.raw)
    case (SubQ16, Some(bv)) => AvmVal.q16(a.raw - bv.raw)
    case (MulQ16, Some(bv)) => AvmVal.q16(q16Mul(a.raw, bv.raw))
    case (DivQ16, Some(bv)) => AvmVal.q16(q16Div(a.raw, bv.raw))
    case (LtQ16, Some(bv))  => AvmVal.bool(a.raw < bv.raw)
    case (EqQ16, Some(bv))  => AvmVal.bool(a.raw == bv.raw)
    case (And, Some(bv))    => AvmVal.bool(a.raw != 0 && bv.raw != 0)
    case (Or, Some(bv))     => AvmVal.bool(a.raw != 0 || bv.raw != 0)
    case (Not, _)           => AvmVal.bool(a.raw == 0)
    case _ => throw new RuntimeException("type mismatch")
  }

  // ── Instructions ────────────────────────────────
  sealed trait Instr
  case class PushQ16(x: Int) extends Instr
  case class PushBool(x: Boolean) extends Instr
  case object Pop extends Instr; case object Dup extends Instr
  case object Swap extends Instr
  case class Load(i: Int) extends Instr; case class Store(i: Int) extends Instr
  case class Jump(t: Int) extends Instr; case class JumpIf(t: Int) extends Instr
  case class PrimOp(p: Prim) extends Instr; case object Halt extends Instr

  // ── State ───────────────────────────────────────
  case class State(pc: Int, stack: List[AvmVal], locals: Vector[Option[AvmVal]],
                   halted: Boolean)
  object State { def apply(nLocals: Int = 0): State =
    State(0, Nil, Vector.fill(nLocals)(None), false)
  }

  // ── Step ────────────────────────────────────────
  def step(s: State, prog: Vector[Instr]): State = {
    if (s.halted) throw new RuntimeException("halted")
    if (s.pc < 0 || s.pc >= prog.length) return s.copy(halted = true)

    prog(s.pc) match {
      case PushQ16(x)  => s.copy(pc = s.pc + 1, stack = AvmVal.q16(x) :: s.stack)
      case PushBool(x) => s.copy(pc = s.pc + 1, stack = AvmVal.bool(x) :: s.stack)
      case Pop         => s.copy(pc = s.pc + 1, stack = s.stack.tail)
      case Dup         => s.copy(pc = s.pc + 1, stack = s.stack.head :: s.stack)
      case Swap        =>
        val a :: b :: rest = s.stack
        s.copy(pc = s.pc + 1, stack = b :: a :: rest)
      case Load(i) =>
        val v = s.locals(i).getOrElse(throw new RuntimeException("missing local"))
        s.copy(pc = s.pc + 1, stack = v :: s.stack)
      case Store(i) =>
        s.copy(pc = s.pc + 1, locals = s.locals.updated(i, Some(s.stack.head)),
               stack = s.stack.tail)
      case Jump(t) => s.copy(pc = t)
      case JumpIf(t) =>
        if (s.stack.head.raw != 0) s.copy(pc = t, stack = s.stack.tail)
        else s.copy(pc = s.pc + 1, stack = s.stack.tail)
      case PrimOp(p) =>
        val arity = if (p == Not) 1 else 2
        val b = if (arity >= 2) Some(s.stack.head) else None
        val rest = if (arity >= 2) s.stack.tail else s.stack
        val a = rest.head
        val result = execPrim(p, a, b)
        s.copy(pc = s.pc + 1, stack = result :: rest.tail)
      case Halt => s.copy(halted = true)
    }
  }

  def run(init: State, prog: Vector[Instr], fuel: Int): State = {
    var s = init
    for (_ <- 0 until fuel if !s.halted) s = step(s, prog)
    s
  }
}