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

package asm

import (
	"fmt"
	"strings"
)

// Encode encodes one Plan 9 instruction (mnemonic plus operands, in source
// order) into x86-64 machine code.
func Encode(mnemonic string, ops ...Operand) ([]byte, error) {
	e := &enc{}
	if err := e.encode(mnemonic, ops); err != nil {
		return nil, err
	}
	return e.out, nil
}

type enc struct {
	out     []byte
	patches []encPatch // disp32 fields awaiting static-symbol resolution
}

// encPatch marks a 4-byte displacement field in enc.out that must receive the
// RIP-relative offset of a static symbol once the file layout is settled.
type encPatch struct {
	off    int
	name   string
	addend int64
}

func (e *enc) encode(mnem string, ops []Operand) error {
	upper := strings.ToUpper(mnem)

	// Fixed-name instructions (no size suffix).
	switch {
	case upper == "RET":
		return e.encodeRet()
	case upper == "NOP":
		return e.emit(&instr{opcode: []byte{0x90}, modrm: -1, sib: -1})
	case upper == "CALL":
		return e.encodeJmpRel(ops, []byte{0xE8})
	case upper == "JMP":
		return e.encodeJmpRel(ops, []byte{0xE9})
	}
	if cc, ok := condCode(upper); ok {
		return e.encodeJcc(cc, ops)
	}

	// VEX (AVX/AVX2) and EVEX (AVX-512) instructions: the trailing
	// B/W/L/Q/D is part of the mnemonic, not a size suffix, so dispatch
	// before splitSize.
	if isVex(upper) || isEvex(upper) || upper == "KMOVW" {
		return e.encodeVec(upper, ops)
	}

	// CMOVcc and SETcc carry the condition in the mnemonic (CMOVLGT, SETNE).
	if strings.HasPrefix(upper, "CMOV") {
		return e.encodeCmov(upper, ops)
	}
	if strings.HasPrefix(upper, "SET") {
		return e.encodeSet(upper, ops)
	}

	base, size := splitSize(upper)
	if size == 0 {
		size = 8 // default operand size in 64-bit mode (e.g. PUSHQ)
	}
	switch base {
	case "MOV":
		return e.encodeMov(ops, size)
	case "ADD", "SUB", "AND", "OR", "XOR", "CMP":
		return e.encodeALU(aluOp[base], ops, size)
	case "TEST":
		return e.encodeTest(ops, size)
	case "LEA":
		return e.encodeLea(ops, size)
	case "INC", "DEC", "NEG", "NOT":
		return e.encodeUnary(unaryOp[base], ops, size)
	case "SHL", "SHR", "SAR":
		return e.encodeShift(shiftOp[base], ops, size)
	case "IMUL", "IMUL3":
		return e.encodeImul(ops, size)
	case "PUSH":
		return e.encodePushPop(ops, true)
	case "POP":
		return e.encodePushPop(ops, false)
	case "LZCNT", "TZCNT":
		return e.encodeCount(base, ops, size)
	case "MOVBLZX", "MOVBQZX", "MOVWLZX", "MOVWQZX", "MOVWLSX", "MOVLQSX":
		return e.encodeMovExtend(base, ops)
	case "CVTSL2SD", "CVTSQ2SD":
		return e.encodeCvtsi2sd(base == "CVTSQ2SD", ops)
	case "MOVOU", "MOVO", "MOVUPS", "MOVAPS", "MOVUPD", "MOVAPD", "MOVSD", "MOVSS":
		return e.encodeSSEMove(sseMoveTable[base], ops)
	}
	return fmt.Errorf("unsupported instruction %q", mnem)
}

// splitSize separates a trailing B/W/L/Q size suffix from the mnemonic.
func splitSize(upper string) (base string, size int) {
	if upper == "" {
		return upper, 0
	}
	switch upper[len(upper)-1] {
	case 'B':
		return upper[:len(upper)-1], 1
	case 'W':
		return upper[:len(upper)-1], 2
	case 'L':
		return upper[:len(upper)-1], 4
	case 'Q':
		return upper[:len(upper)-1], 8
	}
	return upper, 0
}

// encodeVec dispatches a VEX/EVEX mnemonic to the right encoding: KMOVW has
// its own direction-dependent opcodes; KTESTW is always VEX; everything else
// takes EVEX when an operand demands it (a ZMM or K register, or an
// EVEX-only mnemonic) and VEX otherwise.
func (e *enc) encodeVec(upper string, ops []Operand) error {
	if upper == "KMOVW" {
		return e.encodeKmovw(ops)
	}
	if upper == "KTESTW" || !evexRequired(upper, ops) {
		return e.encodeVex(upper, ops)
	}
	return e.encodeEvex(upper, ops)
}

// --- instruction components -------------------------------------------------

type instr struct {
	opSize16  bool
	rexW      bool
	rexR      bool
	rexX      bool
	rexB      bool
	rexForced bool // REX needed even with all bits zero (8-bit low registers)
	prefix    byte // legacy 0xF2/0xF3 prefix (0 = none); emitted after 0x66
	opcode    []byte
	modrm     int // -1 if absent
	sib       int // -1 if absent
	disp      []byte
	imm       []byte
	sb        *sbRef // static-symbol displacement in disp, awaiting resolution
}

// sbRef records that an instruction's displacement refers to a static symbol
// rather than holding a literal value.
type sbRef struct {
	name   string
	addend int64
}

func (e *enc) emit(i *instr) error {
	if i.opSize16 {
		e.out = append(e.out, 0x66)
	}
	if i.prefix != 0 {
		e.out = append(e.out, i.prefix)
	}
	rex := byte(0)
	if i.rexW {
		rex |= 0x08
	}
	if i.rexR {
		rex |= 0x04
	}
	if i.rexX {
		rex |= 0x02
	}
	if i.rexB {
		rex |= 0x01
	}
	if rex != 0 || i.rexForced {
		e.out = append(e.out, 0x40|rex)
	}
	e.out = append(e.out, i.opcode...)
	if i.modrm >= 0 {
		e.out = append(e.out, byte(i.modrm))
	}
	if i.sib >= 0 {
		e.out = append(e.out, byte(i.sib))
	}
	if i.sb != nil {
		e.patches = append(e.patches, encPatch{off: len(e.out), name: i.sb.name, addend: i.sb.addend})
	}
	e.out = append(e.out, i.disp...)
	e.out = append(e.out, i.imm...)
	return nil
}

// newInstr starts an instruction with a size-derived REX.W and 0x66 prefix.
func newInstr(opSize int, opcode []byte) *instr {
	return &instr{
		opSize16: opSize == 2,
		rexW:     opSize == 8,
		opcode:   opcode,
		modrm:    -1,
		sib:      -1,
	}
}

// --- ModR/M, SIB, displacement ----------------------------------------------

// setRM fills in the ModR/M (and SIB, displacement, REX bits) for an
// instruction whose reg field holds a real register `reg` and whose r/m field
// holds `rm`.
func setRM(i *instr, reg Reg, rm Operand, opSize int) error {
	return setRMReg(i, reg.idx&7, reg.idx >= 8, reg.needsREX(opSize), rm, opSize)
}

// setRMDigit fills in the ModR/M for an instruction whose reg field is an
// opcode /digit extension (0–7), which carries none of the register REX rules.
func setRMDigit(i *instr, digit int, rm Operand, opSize int) error {
	return setRMReg(i, digit, false, false, rm, opSize)
}

func setRMReg(i *instr, regField int, rexR, regForced bool, rm Operand, opSize int) error {
	i.rexR = rexR
	if regForced {
		i.rexForced = true
	}

	switch r := rm.(type) {
	case Reg:
		i.rexB = r.idx >= 8
		if r.needsREX(opSize) {
			i.rexForced = true
		}
		i.modrm = 0xC0 | regField<<3 | (r.idx & 7)
		return nil
	case Mem:
		return setMem(i, regField, r)
	case sbMem:
		// RIP-relative reference; the displacement is patched once the static
		// symbol's address is known.
		i.modrm = regField<<3 | 0x05 // mod=00, rm=101 → (RIP)+disp32
		i.disp = le32(0)
		i.sb = &sbRef{name: r.name, addend: r.addend}
		return nil
	default:
		return fmt.Errorf("invalid r/m operand %T", rm)
	}
}

func setMem(i *instr, regField int, m Mem) error {
	modrm, sib, disp, xBit, bBit, err := memComponents(regField, m)
	if err != nil {
		return err
	}
	i.modrm = modrm
	i.sib = sib
	i.disp = disp
	i.rexX = xBit == 1
	i.rexB = bBit == 1
	return nil
}

// memComponents computes the ModR/M byte (with the given reg field), the SIB
// byte (-1 if none), the displacement bytes, and the high index/base bits, for
// a memory operand.  It is shared by the REX (scalar) and VEX (vector) paths.
func memComponents(regField int, m Mem) (modrm, sib int, disp []byte, xBit, bBit int, err error) {
	sib = -1
	// RIP-relative: neither base nor index.
	if !m.HasBase && !m.HasIndex {
		return regField<<3 | 0x05, -1, le32(m.Disp), 0, 0, nil // mod=00, rm=101
	}

	needSIB := m.HasIndex || (m.HasBase && m.Base.idx&7 == 4)

	var mod int
	switch {
	case !m.HasBase:
		mod = 0
		disp = le32(m.Disp)
	case m.Base.idx&7 == 5 && m.Disp == 0:
		mod = 1
		disp = []byte{0}
	case m.Disp == 0:
		mod = 0
	case fits8(m.Disp):
		mod = 1
		disp = []byte{byte(int8(m.Disp))}
	default:
		mod = 2
		disp = le32(m.Disp)
	}

	if needSIB {
		idxField := 4 // 100 = no index
		if m.HasIndex {
			idxField = m.Index.idx & 7
			if m.Index.idx >= 8 {
				xBit = 1
			}
		}
		baseField := 5 // 101 = no base (with mod=00 → disp32)
		if m.HasBase {
			baseField = m.Base.idx & 7
			if m.Base.idx >= 8 {
				bBit = 1
			}
		}
		return mod<<6 | regField<<3 | 0x04, scaleBits(m.Scale)<<6 | idxField<<3 | baseField, disp, xBit, bBit, nil
	}

	if m.Base.idx >= 8 {
		bBit = 1
	}
	return mod<<6 | regField<<3 | (m.Base.idx & 7), -1, disp, 0, bBit, nil
}

func scaleBits(scale int) int {
	switch scale {
	case 2:
		return 1
	case 4:
		return 2
	case 8:
		return 3
	default:
		return 0 // scale 1 (or unset)
	}
}

func fits8(v int64) bool { return v >= -128 && v <= 127 }

func le32(v int64) []byte {
	u := uint32(v)
	return []byte{byte(u), byte(u >> 8), byte(u >> 16), byte(u >> 24)}
}

func le16(v int64) []byte {
	u := uint16(v)
	return []byte{byte(u), byte(u >> 8)}
}

func le64(v int64) []byte {
	u := uint64(v)
	b := make([]byte, 8)
	for i := 0; i < 8; i++ {
		b[i] = byte(u >> (8 * i))
	}
	return b
}