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

// The assembler's side of the amd64 extension layer: this file turns a parsed
// amd64 statement into the operand form arch.ExtInstr.Encode consumes and
// routes statements only the layer can encode through the registry.  It sits
// beside the main amd64 encoders, never inside them: the generated table, the
// legacy SSE paths and the VEX and EVEX mechanisms are untouched, and a
// statement reaches this file only when the mnemonic is registered in the
// extension layer and the scalar paths cannot encode it.
//
// The spellings are the layer's own Plan 9 forms, the ones its metadata
// documents: the vector registers carry the house names X0, Y0 and Z0 (the
// EVEX 128, 256 and 512-bit classes, registers 16 to 31 included), the general
// registers the width their spelling fixes (RAX through R15, EAX through EDI
// and R8D through R15D), the opmask registers K0 through K7, and the memory
// operand the base-relative form off(base) with the optional scaled index
// off(base)(index*scale) the SIB byte carries.  The decorations ride the
// operand in braces: the write mask {k1} through {k7} and zeroing {z} on the
// destination, the {1toN} broadcast on the memory source, and {sae} and
// {rn-sae} through {rz-sae} beside the rounding-capable destinations.  The
// imm8-control forms take their control byte as the leading $ immediate the
// reference listings write first.
//
// The Feature field stays metadata at assembly time: the assembler has no CPU,
// the toolchain does not gate assembly on CPU features, and every registered
// feature assembles, the behaviour the arm64 wiring established
// (asm/arm64_ext.go).  The field remains for the linter and the listing.

package asm

import (
	"fmt"
	"strconv"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-sdk/arch"
	"sourcedock.dev/petrbalvin/gasm-sdk/ast"
)

// amd64ExtStatement converts one instruction's operands into the extended
// layer's operand form.  pinned reports that the statement belongs to the
// layer: the mnemonic is registered in the amd64 registry and the scalar
// paths cannot encode it.  A pinned statement can only encode through the
// layer, so every operand is read here and its diagnostic replaces whatever
// the scalar paths would have said about operands they cannot read; err is
// non-nil for a pinned statement whose operands the layer refuses, and
// extops is complete only when err is nil.  Unpinned means the statement is
// nobody's: the caller falls through to the ordinary amd64 encoders, which
// keep their exact behaviour for every statement they knew before.
func amd64ExtStatement(mnem string, ops []*ast.Operand) (extops []arch.ExtOperand, pinned bool, err error) {
	if _, ok := LookupExtension(arch.AMD64, mnem); !ok {
		return nil, false, nil
	}
	if Encodable(mnem) {
		// A mnemonic the main encoder knows is never the layer's, whatever
		// the registry carries: the scalar paths keep the statement.  No
		// registered mnemonic trips this today (the layer is sealed by
		// test), but the guard keeps the fall-through promise exact should
		// the toolchain ever learn one of these names.
		return nil, false, nil
	}
	out := make([]arch.ExtOperand, 0, len(ops))
	for i, op := range ops {
		ext, convErr := amd64ExtOperand(mnem, op, i+1)
		if convErr != nil {
			return nil, true, convErr
		}
		out = append(out, ext)
	}
	return out, true, nil
}

// amd64ExtOperand converts one parsed operand into the layer's form: a $ immediate,
// a vector, general or opmask register, or a base-relative memory operand, each
// with the brace decorations the spelling carries.
func amd64ExtOperand(mnem string, op *ast.Operand, pos int) (arch.ExtOperand, error) {
	if op.Kind == ast.OpImmediate {
		return amd64ExtImmediate(mnem, op, pos)
	}
	body, dec, err := amd64ExtDecorations(mnem, op, pos)
	if err != nil {
		return arch.ExtOperand{}, err
	}
	if strings.ContainsRune(body, '(') {
		ext, ok := amd64ExtMemory(mnem, op, pos)
		if !ok {
			return arch.ExtOperand{}, fmt.Errorf("%s: operand %d (%s) is not an extended-layer operand: want a base-relative memory operand, off(base)(index*scale) shape", mnem, pos, op.Raw)
		}
		ext.Broadcast = dec.broadcast
		if dec.hasMask {
			ext.Mask, ext.HasMask = dec.mask, true
		}
		ext.Zeroing = dec.zeroing
		ext.Round = dec.round
		return ext, nil
	}
	ext, ok := amd64ExtRegister(mnem, op, pos, body)
	if !ok {
		return arch.ExtOperand{}, fmt.Errorf("%s: operand %d (%s) is not an extended-layer operand: want a vector, general or opmask register, a base-relative memory operand or an immediate", mnem, pos, op.Raw)
	}
	if dec.hasMask {
		ext.Mask, ext.HasMask = dec.mask, true
	}
	ext.Zeroing = dec.zeroing
	ext.Round = dec.round
	return ext, nil
}

// amd64ExtImmediate converts a $ immediate into the layer's form.  The
// parser folds a parenthesised constant expression in full and reads a bare
// literal greedily, dropping any trailing operator tokens: $255<<8 parses as
// 255 with the shift silently gone.  Encoding that silent prefix would
// assemble what the text did not say, so an unparenthesised immediate is
// accepted only when its whole text reads back as one integer carrying the
// parser's value.
func amd64ExtImmediate(mnem string, op *ast.Operand, pos int) (arch.ExtOperand, error) {
	if !op.Imm.HasVal {
		return arch.ExtOperand{}, fmt.Errorf("%s: operand %d (%s) is not an immediate the layer can read", mnem, pos, op.Raw)
	}
	text := strings.Join(strings.Fields(strings.TrimPrefix(op.Raw, "$")), "")
	if !strings.HasPrefix(text, "(") {
		if _, parseErr := strconv.ParseInt(text, 0, 64); parseErr != nil {
			return arch.ExtOperand{}, fmt.Errorf("%s: operand %d (%s) is not an immediate the layer can read", mnem, pos, op.Raw)
		}
	}
	v := op.Imm.Val
	if op.Imm.Neg {
		v = -v
	}
	return arch.ExtOperand{Kind: arch.ExtImm, Imm: v}, nil
}

// amd64ExtRegister parses a register operand off a normalised operand body:
// the vector classes X0-X31, Y0-Y31 and Z0-Z31, the width-fixed general
// spellings RAX through R15 and EAX through R15D, and the opmask registers
// K0-K7.  The register ranges are left to the encoding, whose diagnostics
// name them.
func amd64ExtRegister(mnem string, op *ast.Operand, pos int, body string) (arch.ExtOperand, bool) {
	if body == "" {
		return arch.ExtOperand{}, false
	}
	r, ok := ParseReg(body)
	if !ok {
		return arch.ExtOperand{}, false
	}
	switch {
	case r.mask:
		return arch.ExtOperand{Kind: arch.ExtKReg, Reg: r.idx}, true
	case r.isVec():
		kind := arch.ExtXMM
		switch r.size {
		case 32:
			kind = arch.ExtYMM
		case 64:
			kind = arch.ExtZMM
		}
		return arch.ExtOperand{Kind: kind, Reg: r.idx}, true
	case r.size == 8:
		return arch.ExtOperand{Kind: arch.ExtR64, Reg: r.idx}, true
	case r.size == 4:
		return arch.ExtOperand{Kind: arch.ExtR32, Reg: r.idx}, true
	}
	return arch.ExtOperand{}, false
}

// amd64ExtMemory parses a base-relative memory operand off the parsed
// address: off(base) and off(base)(index*scale), the SIB shapes the layer's
// entries carry.  The base and the index are general registers spelled in any
// width the house names offer, the displacement the leading signed term, and
// a group whose scale is not written scales by one, the choice the main
// amd64 paths make for the same spelling.  Vector, opmask and segment
// registers are refused as base and index, and so is every frame form: the
// layer's memory operand is hardware addressing alone.
func amd64ExtMemory(mnem string, op *ast.Operand, pos int) (arch.ExtOperand, bool) {
	a := op.Addr
	if a.Range != nil || a.Base == "" {
		return arch.ExtOperand{}, false
	}
	if a.Sym != nil && a.Sym.Pseudo != "" {
		return arch.ExtOperand{}, false
	}
	base, ok := amd64ExtGprNumber(a.Base)
	if !ok {
		return arch.ExtOperand{}, false
	}
	ext := arch.ExtOperand{Kind: arch.ExtMem, Reg: base, Imm: a.Offset}
	if a.Index != "" {
		index, ok := amd64ExtGprNumber(a.Index)
		if !ok {
			return arch.ExtOperand{}, false
		}
		scale := a.Scale
		if scale == 0 {
			scale = 1
		}
		ext.Index, ext.Scale, ext.HasIndex = index, scale, true
	}
	return ext, true
}

// amd64ExtGprNumber resolves one general-register spelling to its number:
// whatever the register table carries for indices 0-15, the vector, opmask,
// x87, MMX, segment and control-debug classes refused, so a vector register
// in a base or index position names itself rather than encoding as its
// same-numbered general register.
func amd64ExtGprNumber(name string) (int, bool) {
	r, ok := ParseReg(name)
	if !ok || r.mask || r.fp || r.mmx || r.seg != 0 || r.ctl != 0 || r.size > 8 {
		return 0, false
	}
	return r.idx, true
}

// amd64ExtDecorations splits the brace decorations off a normalised operand
// text and returns the body before the first brace and the decorations they
// spell: the write mask {k1} through {k7}, zeroing {z}, the {1toN} broadcast
// and the rounding controls {sae} and {rn-sae} through {rz-sae}, matched
// case-insensitively the way the register spellings are.  The mask, zeroing
// and rounding fields land on the operand the conversion builds; whether the
// position takes them is the encoding's judgement, whose diagnostics name the
// entry.  The broadcast factor N is checked as a number and otherwise left to
// the entry: the layer's model carries the spelling, not the lane count.
func amd64ExtDecorations(mnem string, op *ast.Operand, pos int) (body string, dec amd64ExtDecor, err error) {
	compact := strings.Join(strings.Fields(op.Raw), "")
	i := strings.IndexByte(compact, '{')
	if i < 0 {
		return compact, dec, nil
	}
	body = compact[:i]
	for i < len(compact) {
		if compact[i] != '{' {
			return "", dec, fmt.Errorf("%s: operand %d (%s): text between brace decorations", mnem, pos, op.Raw)
		}
		end := strings.IndexByte(compact[i:], '}')
		if end < 0 {
			return "", dec, fmt.Errorf("%s: operand %d (%s): brace decoration without a closing brace", mnem, pos, op.Raw)
		}
		content := strings.ToUpper(compact[i+1 : i+end])
		switch {
		case content == "Z":
			if dec.zeroing {
				return "", dec, fmt.Errorf("%s: operand %d (%s) carries two zeroing decorations", mnem, pos, op.Raw)
			}
			dec.zeroing = true
		case content == "SAE":
			if dec.round != arch.ExtRoundNone {
				return "", dec, fmt.Errorf("%s: operand %d (%s) carries two rounding controls", mnem, pos, op.Raw)
			}
			dec.round = arch.ExtRoundSAE
		case content == "RN-SAE":
			if dec.round != arch.ExtRoundNone {
				return "", dec, fmt.Errorf("%s: operand %d (%s) carries two rounding controls", mnem, pos, op.Raw)
			}
			dec.round = arch.ExtRoundNearest
		case content == "RD-SAE":
			if dec.round != arch.ExtRoundNone {
				return "", dec, fmt.Errorf("%s: operand %d (%s) carries two rounding controls", mnem, pos, op.Raw)
			}
			dec.round = arch.ExtRoundDown
		case content == "RU-SAE":
			if dec.round != arch.ExtRoundNone {
				return "", dec, fmt.Errorf("%s: operand %d (%s) carries two rounding controls", mnem, pos, op.Raw)
			}
			dec.round = arch.ExtRoundUp
		case content == "RZ-SAE":
			if dec.round != arch.ExtRoundNone {
				return "", dec, fmt.Errorf("%s: operand %d (%s) carries two rounding controls", mnem, pos, op.Raw)
			}
			dec.round = arch.ExtRoundTruncate
		case strings.HasPrefix(content, "K") && content != "K":
			n, convErr := strconv.Atoi(content[1:])
			if convErr != nil || n < 0 {
				return "", dec, fmt.Errorf("%s: operand %d (%s): %q is not a mask decoration, want {k1} through {k7}", mnem, pos, op.Raw, content)
			}
			if dec.hasMask {
				return "", dec, fmt.Errorf("%s: operand %d (%s) carries two write masks", mnem, pos, op.Raw)
			}
			dec.mask, dec.hasMask = n, true
		case strings.HasPrefix(content, "1TO"):
			if _, convErr := strconv.Atoi(content[3:]); convErr != nil {
				return "", dec, fmt.Errorf("%s: operand %d (%s): %q is not a broadcast decoration, want {1toN}", mnem, pos, op.Raw, content)
			}
			dec.broadcast = true
		default:
			return "", dec, fmt.Errorf("%s: operand %d (%s): {%s} is not a decoration the layer reads: want {k1} through {k7}, {z}, {1toN}, {sae} or {rn-sae} through {rz-sae}", mnem, pos, op.Raw, content)
		}
		i += end + 1
	}
	return body, dec, nil
}

// amd64ExtDecor carries the brace decorations one operand's spelling names.
type amd64ExtDecor struct {
	mask      int
	hasMask   bool
	zeroing   bool
	broadcast bool
	round     arch.ExtRounding
}