// Package vm implements the bytecode CPU that runs ship programs.
//
// The machine has 16 32-bit registers (r15 is the stack pointer), a program
// ROM and a byte-addressable little-endian data RAM. Peripherals are reached
// through IN/OUT port instructions. Every instruction is 4 bytes:
//
//	byte 0: opcode | byte 1: ra<<4 | rb | bytes 2-3: signed 16-bit immediate
//
// Branch/jump immediates are offsets in instructions relative to the next
// instruction. Each tick a CPU runs up to a cycle budget or until it executes
// YIELD; state persists between ticks so an over-long computation simply
// continues on the next tick.
package vm

import (
	"encoding/binary"
	"fmt"
)

type Op uint8

const (
	NOP Op = iota
	YIELD
	HALT
	LDI
	LUI
	MOV
	ADD
	SUB
	MUL
	DIV
	MOD
	AND
	OR
	XOR
	SHL
	SHR
	SAR
	ADDI
	JMP
	BEQ
	BNE
	BLT
	BGE
	CALL
	RET
	PUSH
	POP
	LDB
	LDH
	LDW
	STB
	STH
	STW
	IN
	OUT
	numOps
)

const (
	NumRegs = 16
	SP      = 15
)

type Status uint8

const (
	Running Status = iota
	Yielded        // finished this tick's work voluntarily
	Halted
	Faulted
)

func (s Status) String() string {
	return [...]string{"running", "yielded", "halted", "faulted"}[s]
}

// Bus connects the CPU to peripherals.
type Bus interface {
	In(port uint16) int32
	Out(port uint16, v int32)
}

type CPU struct {
	R      [NumRegs]int32
	PC     uint32 // byte address into Prog
	Prog   []byte
	RAM    []byte
	Status Status
	Fault  string
}

// New creates a CPU with the program loaded and the stack pointer at the top
// of RAM.
func New(prog []byte, ramBytes int) (*CPU, error) {
	if len(prog)%4 != 0 {
		return nil, fmt.Errorf("program length %d is not a multiple of 4", len(prog))
	}
	c := &CPU{Prog: append([]byte(nil), prog...), RAM: make([]byte, ramBytes)}
	c.R[SP] = int32(ramBytes)
	return c, nil
}

func (c *CPU) fault(format string, args ...any) {
	c.Status = Faulted
	c.Fault = fmt.Sprintf(format, args...)
}

func cost(op Op) int {
	switch op {
	case MUL:
		return 2
	case DIV, MOD:
		return 8
	}
	return 1
}

// Run executes instructions until the budget is spent, YIELD, HALT or a
// fault, returning the cycles used. A yielded CPU resumes on the next call.
func (c *CPU) Run(bus Bus, budget int) int {
	if c.Status == Halted || c.Status == Faulted {
		return 0
	}
	c.Status = Running
	used := 0
	for used < budget {
		if int(c.PC)+4 > len(c.Prog) {
			// Falling off the end of the program halts the CPU.
			c.Status = Halted
			break
		}
		w := binary.LittleEndian.Uint32(c.Prog[c.PC:])
		op := Op(w)
		ra, rb := (w>>12)&0xf, (w>>8)&0xf
		imm := int32(int16(w >> 16))
		if op >= numOps {
			c.fault("illegal opcode %d at pc=%d", op, c.PC)
			break
		}
		used += cost(op)
		next := c.PC + 4
		r := &c.R
		switch op {
		case NOP:
		case YIELD:
			c.Status = Yielded
		case HALT:
			c.Status = Halted
		case LDI:
			r[ra] = imm
		case LUI:
			r[ra] = int32(uint32(imm)<<16 | uint32(r[ra])&0xffff)
		case MOV:
			r[ra] = r[rb]
		case ADD:
			r[ra] += r[rb]
		case SUB:
			r[ra] -= r[rb]
		case MUL:
			r[ra] *= r[rb]
		case DIV, MOD:
			d := r[rb]
			if d == 0 {
				c.fault("division by zero at pc=%d", c.PC)
				break
			}
			switch {
			case d == -1: // avoid MinInt32 / -1 overflow panic semantics
				if op == DIV {
					r[ra] = -r[ra]
				} else {
					r[ra] = 0
				}
			case op == DIV:
				r[ra] /= d
			default:
				r[ra] %= d
			}
		case AND:
			r[ra] &= r[rb]
		case OR:
			r[ra] |= r[rb]
		case XOR:
			r[ra] ^= r[rb]
		case SHL:
			r[ra] = int32(uint32(r[ra]) << (uint32(r[rb]) & 31))
		case SHR:
			r[ra] = int32(uint32(r[ra]) >> (uint32(r[rb]) & 31))
		case SAR:
			r[ra] >>= uint32(r[rb]) & 31
		case ADDI:
			r[ra] += imm
		case JMP:
			next = uint32(int64(next) + int64(imm)*4)
		case BEQ, BNE, BLT, BGE:
			var t bool
			switch op {
			case BEQ:
				t = r[ra] == r[rb]
			case BNE:
				t = r[ra] != r[rb]
			case BLT:
				t = r[ra] < r[rb]
			case BGE:
				t = r[ra] >= r[rb]
			}
			if t {
				next = uint32(int64(next) + int64(imm)*4)
			}
		case CALL:
			if c.push(int32(next)) {
				next = uint32(int64(next) + int64(imm)*4)
			}
		case RET:
			if v, ok := c.pop(); ok {
				next = uint32(v)
			}
		case PUSH:
			c.push(r[ra])
		case POP:
			if v, ok := c.pop(); ok {
				r[ra] = v
			}
		case LDB, LDH, LDW:
			n := accessSize(op)
			if a, ok := c.addr(r[rb]+imm, n); ok {
				var v uint32
				for i := n - 1; i >= 0; i-- {
					v = v<<8 | uint32(c.RAM[a+i])
				}
				r[ra] = int32(v)
			}
		case STB, STH, STW:
			n := accessSize(op)
			if a, ok := c.addr(r[rb]+imm, n); ok {
				v := uint32(r[ra])
				for i := 0; i < n; i++ {
					c.RAM[a+i] = byte(v >> (8 * i))
				}
			}
		case IN:
			r[ra] = bus.In(uint16(imm))
		case OUT:
			bus.Out(uint16(imm), r[ra])
		}
		if c.Status == Faulted {
			break
		}
		c.PC = next
		if c.Status != Running {
			break
		}
	}
	return used
}

func accessSize(op Op) int {
	switch op {
	case LDB, STB:
		return 1
	case LDH, STH:
		return 2
	}
	return 4
}

func (c *CPU) addr(a int32, n int) (int, bool) {
	if a < 0 || int(a)+n > len(c.RAM) {
		c.fault("memory access out of range: %d", a)
		return 0, false
	}
	return int(a), true
}

func (c *CPU) push(v int32) bool {
	a, ok := c.addr(c.R[SP]-4, 4)
	if !ok {
		return false
	}
	c.R[SP] -= 4
	binary.LittleEndian.PutUint32(c.RAM[a:], uint32(v))
	return true
}

func (c *CPU) pop() (int32, bool) {
	a, ok := c.addr(c.R[SP], 4)
	if !ok {
		return 0, false
	}
	c.R[SP] += 4
	return int32(binary.LittleEndian.Uint32(c.RAM[a:])), true
}

// Encode builds one instruction word.
func Encode(op Op, ra, rb int, imm int32) uint32 {
	return uint32(op) | uint32(rb&0xf)<<8 | uint32(ra&0xf)<<12 | uint32(uint16(imm))<<16
}

// MarshalState serialises the mutable CPU state (not the program).
func (c *CPU) MarshalState() []byte {
	b := make([]byte, 0, NumRegs*4+8+len(c.RAM))
	for _, r := range c.R {
		b = binary.LittleEndian.AppendUint32(b, uint32(r))
	}
	b = binary.LittleEndian.AppendUint32(b, c.PC)
	b = append(b, byte(c.Status), 0, 0, 0)
	return append(b, c.RAM...)
}

// UnmarshalState restores state produced by MarshalState.
func (c *CPU) UnmarshalState(b []byte) error {
	hdr := NumRegs*4 + 8
	if len(b) != hdr+len(c.RAM) {
		return fmt.Errorf("state size %d does not match expected %d", len(b), hdr+len(c.RAM))
	}
	for i := range c.R {
		c.R[i] = int32(binary.LittleEndian.Uint32(b[i*4:]))
	}
	c.PC = binary.LittleEndian.Uint32(b[NumRegs*4:])
	c.Status = Status(b[NumRegs*4+4])
	copy(c.RAM, b[hdr:])
	return nil
}