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

import "errors"

// ── Constants ────────────────────────────────────────────────────
const (
	AVMClampMin = -2147483647
	AVMClampMax = 2147483647
	AVMQ0Min    = -32767
	AVMQ0Max    = 32767
	Q16Scale    = 65536
	AVMMaxStack = 1024
)

func avmClamp(x int64) int32 {
	if x > AVMClampMax { return AVMClampMax }
	if x < AVMClampMin { return AVMClampMin }
	return int32(x)
}

func avmQ0Clamp(x int64) int32 {
	if x > AVMQ0Max { return AVMQ0Max }
	if x < AVMQ0Min { return AVMQ0Min }
	return int32(x)
}

func floorDiv(a, b int64) int64 {
	if b == 0 { return 0 }
	q := a / b
	r := a % b
	if r != 0 && ((a ^ b) < 0) { q-- }
	return q
}

// euclideanDiv matches Lean Int.ediv (remainder >= 0).
func euclideanDiv(a, b int64) int64 {
	if b == 0 { return 0 }
	q := a / b
	r := a % b
	if r < 0 {
		if b > 0 { q-- } else { q++ }
	}
	return q
}

func ltQ16V6(a, b int32) bool {
	sa, sb := a < 0, b < 0
	if sa != sb { return sa }
	return a < b
}

// ── Types ──────────────────────────────────────────────────────
type Ty uint8
const (
	Q0_16  Ty = 0
	Q16_16 Ty = 1
	Bool   Ty = 2
)

type Val struct {
	Ty   Ty
	Q    int32
	Bval bool
}

// ── Primitives ────────────────────────────────────────────────
type Prim uint8
const (
	AddSatQ0 Prim = iota; SubSatQ0; AddSatQ16; SubSatQ16
	MulSatQ16; DivSatQ16; LtQ16; EqQ16
	And; Or; Not
)

func primArity(p Prim) int {
	if p == Not { return 1 }
	return 2
}

// ── Instructions ──────────────────────────────────────────────
type Op uint8
const (
	PushQ16 Op = iota; PushBool; PushQ0; Pop; Dup; Swap
	Load; Store; Jump; JumpIf; Primitive; Halt
)

type Instr struct {
	Op   Op
	Arg  int32
	Arg2 bool
}

// ── Primitive execution ──────────────────────────────────────
func execPrim(p Prim, a, b Val) (Val, error) {
	check := func(v Val, t Ty) error {
		if v.Ty != t { return errors.New("type mismatch") }
		return nil
	}
	switch p {
	case AddSatQ0:
		if err := check(a, Q0_16); err != nil { return Val{}, err }
		if err := check(b, Q0_16); err != nil { return Val{}, err }
		return Val{Ty: Q0_16, Q: avmQ0Clamp(int64(a.Q) + int64(b.Q))}, nil
	case SubSatQ0:
		if err := check(a, Q0_16); err != nil { return Val{}, err }
		if err := check(b, Q0_16); err != nil { return Val{}, err }
		return Val{Ty: Q0_16, Q: avmQ0Clamp(int64(a.Q) - int64(b.Q))}, nil
	case AddSatQ16:
		if err := check(a, Q16_16); err != nil { return Val{}, err }
		if err := check(b, Q16_16); err != nil { return Val{}, err }
		return Val{Ty: Q16_16, Q: avmClamp(int64(a.Q) + int64(b.Q))}, nil
	case SubSatQ16:
		if err := check(a, Q16_16); err != nil { return Val{}, err }
		if err := check(b, Q16_16); err != nil { return Val{}, err }
		return Val{Ty: Q16_16, Q: avmClamp(int64(a.Q) - int64(b.Q))}, nil
	case MulSatQ16:
		if err := check(a, Q16_16); err != nil { return Val{}, err }
		if err := check(b, Q16_16); err != nil { return Val{}, err }
		return Val{Ty: Q16_16, Q: avmClamp(euclideanDiv(int64(a.Q)*int64(b.Q), Q16Scale))}, nil
	case DivSatQ16:
		if err := check(a, Q16_16); err != nil { return Val{}, err }
		if err := check(b, Q16_16); err != nil { return Val{}, err }
		if b.Q == 0 { return Val{}, errors.New("division by zero") }
		return Val{Ty: Q16_16, Q: avmClamp(euclideanDiv(int64(a.Q)*Q16Scale, int64(b.Q)))}, nil
	case LtQ16:
		if err := check(a, Q16_16); err != nil { return Val{}, err }
		if err := check(b, Q16_16); err != nil { return Val{}, err }
		return Val{Ty: Bool, Bval: ltQ16V6(a.Q, b.Q)}, nil
	case EqQ16:
		if err := check(a, Q16_16); err != nil { return Val{}, err }
		if err := check(b, Q16_16); err != nil { return Val{}, err }
		return Val{Ty: Bool, Bval: a.Q == b.Q}, nil
	case And:
		if err := check(a, Bool); err != nil { return Val{}, err }
		if err := check(b, Bool); err != nil { return Val{}, err }
		return Val{Ty: Bool, Bval: a.Bval && b.Bval}, nil
	case Or:
		if err := check(a, Bool); err != nil { return Val{}, err }
		if err := check(b, Bool); err != nil { return Val{}, err }
		return Val{Ty: Bool, Bval: a.Bval || b.Bval}, nil
	case Not:
		if err := check(a, Bool); err != nil { return Val{}, err }
		return Val{Ty: Bool, Bval: !a.Bval}, nil
	}
	return Val{}, errors.New("unknown prim")
}

// ── State ─────────────────────────────────────────────────────
type State struct {
	Pc     int
	Stack  []Val
	Locals []Val
	Halted bool
}

func NewState(nLocals int) State {
	return State{Pc: 0, Stack: []Val{}, Locals: make([]Val, nLocals), Halted: false}
}

// ── Step ─────────────────────────────────────────────────────
func Step(s State, prog []Instr) (*State, error) {
	if s.Halted { return &s, nil }
	if s.Pc < 0 || s.Pc >= len(prog) {
		return nil, errors.New("invalid jump")
	}

	instr := prog[s.Pc]
	stack := s.Stack
	npc := s.Pc + 1

	growing := instr.Op == PushQ16 || instr.Op == PushBool || instr.Op == PushQ0 || instr.Op == Dup || instr.Op == Load
	if growing && len(stack) >= AVMMaxStack { return nil, errors.New("stack overflow") }

	switch instr.Op {
	case PushQ16:
		stack = append(stack, Val{Ty: Q16_16, Q: avmClamp(int64(instr.Arg))})
		case PushBool:
			stack = append(stack, Val{Ty: Bool, Bval: instr.Arg2})
		case PushQ0:
			stack = append(stack, Val{Ty: Q0_16, Q: avmQ0Clamp(int64(instr.Arg))})
	case Pop:
		if len(stack) == 0 { return nil, errors.New("empty stack") }
		stack = stack[:len(stack)-1]
	case Dup:
		if len(stack) == 0 { return nil, errors.New("empty stack") }
		stack = append(stack, stack[len(stack)-1])
	case Swap:
		if len(stack) < 2 { return nil, errors.New("stack underflow") }
		stack[len(stack)-1], stack[len(stack)-2] = stack[len(stack)-2], stack[len(stack)-1]
	case Load:
		i := int(instr.Arg)
		if i >= len(s.Locals) { return nil, errors.New("missing local") }
		stack = append(stack, s.Locals[i])
	case Store:
		if len(stack) == 0 { return nil, errors.New("empty stack") }
		i := int(instr.Arg)
		if i >= len(s.Locals) { return nil, errors.New("missing local") }
		s.Locals[i] = stack[len(stack)-1]
		stack = stack[:len(stack)-1]
	case Jump:
		if instr.Arg < 0 || int(instr.Arg) >= len(prog) { return nil, errors.New("jump OOB") }
		npc = int(instr.Arg)
	case JumpIf:
		if len(stack) == 0 { return nil, errors.New("empty stack") }
		v := stack[len(stack)-1]; stack = stack[:len(stack)-1]
		if v.Ty != Bool { return nil, errors.New("type mismatch") }
		if v.Bval {
			if instr.Arg < 0 || int(instr.Arg) >= len(prog) { return nil, errors.New("jump OOB") }
			npc = int(instr.Arg)
		}
	case Primitive:
		p := Prim(instr.Arg)
		arity := primArity(p)
		if len(stack) < arity { return nil, errors.New("stack underflow") }
		var b Val
		if arity >= 2 { b = stack[len(stack)-1]; stack = stack[:len(stack)-1] }
		a := stack[len(stack)-1]; stack = stack[:len(stack)-1]
		r, execErr := execPrim(p, a, b)
		if execErr != nil { return nil, execErr }
		stack = append(stack, r)
	case Halt:
		s.Halted = true
		npc = s.Pc  // Lean: Halt does not increment PC
	}
	return &State{Pc: npc, Stack: stack, Locals: s.Locals, Halted: s.Halted}, nil
}

// ── Run (fuel-bounded) ───────────────────────────────────────
func Run(init State, prog []Instr, fuel int) (*State, error) {
	s := init
	for i := 0; i < fuel; i++ {
		if s.Halted { return &s, nil }
		next, err := Step(s, prog)
		if err != nil { return nil, err }
		s = *next
	}
	return &s, nil
}