// Copyright (c) 2026 Petr Balvín <opensource@petrbalvin.org> (https://petrbalvin.org)
// SPDX-License-Identifier: BSD-3-Clause

// Package asm is a standalone assembler: it encodes Plan 9 assembly
// instructions into machine code without the Go toolchain.  Phase 2 begins
// with an amd64 (x86-64) scalar instruction encoder; the encoding is validated
// by round-tripping through golang.org/x/arch's decoder in the tests.
package asm

import "maps"

import "strings"

// Reg is an x86-64 register.  In Plan 9 assembly the classic names (AX, BX, …)
// are size-agnostic — the instruction suffix (MOVQ vs MOVL) fixes the width —
// so the encoder keys off the register's index and lets the mnemonic supply the
// size.  The high flag marks the legacy high-byte registers AH/CH/DH/BH, which
// occupy indices 4–7 yet take no REX prefix, unlike SPL/BPL/SIL/DIL that share
// those indices but require one.  The mask flag marks the AVX-512 opmask
// registers K0–K7.
type Reg struct {
	idx  int
	size int  // informational width implied by the name; the mnemonic decides
	high bool // AH/CH/DH/BH
	mask bool // K0–K7 opmask register
}

// Index returns the register number (0–15 for GPRs, 0–31 for vectors).
func (r Reg) Index() int { return r.idx }

// Size returns the width in bytes implied by the register's name.
func (r Reg) Size() int { return r.size }

// IsMask reports whether r is an AVX-512 opmask register (K0–K7).
func (r Reg) IsMask() bool { return r.mask }

func (r Reg) isOperand() {}

// needsREX reports whether this register forces a REX prefix at the given
// operand size: the extended registers R8–R15 always do, and at byte size the
// low registers SPL/BPL/SIL/DIL (indices 4–7, not high) do as well.
func (r Reg) needsREX(opSize int) bool {
	if r.idx >= 8 {
		return true
	}
	return opSize == 1 && r.idx >= 4 && !r.high
}

// Register constants (the size is the width the name implies).
var (
	AL  = Reg{idx: 0, size: 1}
	CL  = Reg{idx: 1, size: 1}
	DL  = Reg{idx: 2, size: 1}
	BL  = Reg{idx: 3, size: 1}
	AH  = Reg{idx: 4, size: 1, high: true}
	CH  = Reg{idx: 5, size: 1, high: true}
	DH  = Reg{idx: 6, size: 1, high: true}
	BH  = Reg{idx: 7, size: 1, high: true}
	SPL = Reg{idx: 4, size: 1}
	BPL = Reg{idx: 5, size: 1}
	SIL = Reg{idx: 6, size: 1}
	DIL = Reg{idx: 7, size: 1}

	AX = Reg{idx: 0, size: 2}
	CX = Reg{idx: 1, size: 2}
	DX = Reg{idx: 2, size: 2}
	BX = Reg{idx: 3, size: 2}
	_  = Reg{idx: 4, size: 2}
	_  = Reg{idx: 5, size: 2}
	SI = Reg{idx: 6, size: 2}
	DI = Reg{idx: 7, size: 2}

	_ = Reg{idx: 0, size: 4}
	_ = Reg{idx: 1, size: 4}
	_ = Reg{idx: 2, size: 4}
	_ = Reg{idx: 3, size: 4}
	_ = Reg{idx: 4, size: 4}
	_ = Reg{idx: 5, size: 4}
	_ = Reg{idx: 6, size: 4}
	_ = Reg{idx: 7, size: 4}

	_ = Reg{idx: 0, size: 8}
	_ = Reg{idx: 1, size: 8}
	_ = Reg{idx: 2, size: 8}
	_ = Reg{idx: 3, size: 8}
	_ = Reg{idx: 4, size: 8}
	_ = Reg{idx: 5, size: 8}
	_ = Reg{idx: 6, size: 8}
	_ = Reg{idx: 7, size: 8}
)

// regByName maps an assembly register name (case-insensitive) to a Reg.
var regByName = buildRegByName()

func buildRegByName() map[string]Reg {
	m := map[string]Reg{}

	// 64-bit: RAX..RDI, R8..R15.
	r64 := []string{"RAX", "RCX", "RDX", "RBX", "RSP", "RBP", "RSI", "RDI"}
	for i, n := range r64 {
		m[n] = Reg{idx: i, size: 8}
	}
	for i := 8; i <= 15; i++ {
		m["R"+itoa(i)] = Reg{idx: i, size: 8}
	}

	// 32-bit: EAX..EDI, R8D..R15D.
	e32 := []string{"EAX", "ECX", "EDX", "EBX", "ESP", "EBP", "ESI", "EDI"}
	for i, n := range e32 {
		m[n] = Reg{idx: i, size: 4}
	}
	for i := 8; i <= 15; i++ {
		m["R"+itoa(i)+"D"] = Reg{idx: i, size: 4}
	}

	// 16-bit: AX..DI, R8W..R15W.
	w16 := []string{"AX", "CX", "DX", "BX", "SP", "BP", "SI", "DI"}
	for i, n := range w16 {
		m[n] = Reg{idx: i, size: 2}
	}
	for i := 8; i <= 15; i++ {
		m["R"+itoa(i)+"W"] = Reg{idx: i, size: 2}
	}

	// 8-bit: AL..BH, SPL..DIL, R8B..R15B.
	maps.Copy(m, map[string]Reg{
		"AL": AL, "CL": CL, "DL": DL, "BL": BL,
		"AH": AH, "CH": CH, "DH": DH, "BH": BH,
		"SPL": SPL, "BPL": BPL, "SIL": SIL, "DIL": DIL,
	})
	for i := 8; i <= 15; i++ {
		m["R"+itoa(i)+"B"] = Reg{idx: i, size: 1}
	}

	// Vector: X0..X31 (128-bit, size 16), Y0..Y31 (256-bit, size 32),
	// Z0..Z31 (512-bit, size 64).  Indices 16–31 are only encodable in EVEX
	// (AVX-512) instructions; the encoder validates that through its tables.
	for i := 0; i <= 31; i++ {
		m["X"+itoa(i)] = Reg{idx: i, size: 16}
		m["Y"+itoa(i)] = Reg{idx: i, size: 32}
		m["Z"+itoa(i)] = Reg{idx: i, size: 64}
	}
	// Opmask: K0..K7.
	for i := 0; i <= 7; i++ {
		m["K"+itoa(i)] = Reg{idx: i, size: 8, mask: true}
	}
	return m
}

// isVec reports whether r is an XMM/YMM/ZMM vector register.
func (r Reg) isVec() bool { return r.size == 16 || r.size == 32 || r.size == 64 }

// vecLenBit returns the vector-length field for a vector register:
// 0 (128-bit, VEX.L / EVEX.L'L=00), 1 (256-bit) or 2 (512-bit, EVEX only).
func (r Reg) vecLenBit() int {
	switch r.size {
	case 64:
		return 2
	case 32:
		return 1
	}
	return 0
}

// ParseReg resolves an assembly register name to a Reg.
func ParseReg(name string) (Reg, bool) {
	r, ok := regByName[strings.ToUpper(name)]
	return r, ok
}

func itoa(n int) string {
	if n == 0 {
		return "0"
	}
	var buf [3]byte
	i := len(buf)
	for n > 0 {
		i--
		buf[i] = byte('0' + n%10)
		n /= 10
	}
	return string(buf[i:])
}