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

package asm

import (
	"encoding/hex"
	"strings"
	"testing"

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

// TestExtensionRegistryARM64 checks the mnemonic lookup over and above the
// generated arm64 table: one mnemonic, several forms, case-insensitive, and
// nothing offered for a spelling the layer does not carry.
func TestExtensionRegistryARM64(t *testing.T) {
	add, ok := LookupExtension(arch.ARM64, "ADD")
	if !ok {
		t.Fatal("LookupExtension(ARM64, ADD) found nothing")
	}
	var forms []arch.ExtForm
	for _, in := range add {
		if in.Name != "ADD" {
			t.Errorf("candidate %q leaked into the ADD lookup", in.Name)
		}
		forms = append(forms, in.Form)
	}
	if len(forms) != 3 ||
		forms[0] != arch.ExtFormVectors ||
		forms[1] != arch.ExtFormPredicated ||
		forms[2] != arch.ExtFormImmediate {
		t.Errorf("ADD registers forms %v, want unpredicated, predicated and immediate", forms)
	}
	if _, ok := LookupExtension(arch.ARM64, "add"); !ok {
		t.Error("the lookup is case-sensitive")
	}
	if _, ok := LookupExtension(arch.ARM64, "NOSUCHINSTR"); ok {
		t.Error("a non-extended mnemonic resolved")
	}
	sqadd, ok := LookupExtension(arch.ARM64, "SQADD")
	if !ok || len(sqadd) != 2 {
		t.Errorf("SQADD registers %d forms, want the unpredicated and immediate pair", len(sqadd))
	}
}

// TestExtensionAboveGeneratedTable pins the layering: SQADD is nowhere in the
// generated arm64 table (the toolchain knows only the NEON spelling VSQADD)
// yet the extension layer carries it, while ADD sits in both layers
// independently.
func TestExtensionAboveGeneratedTable(t *testing.T) {
	if _, found := arch.ForArch(arch.ARM64).Lookup("SQADD"); found {
		t.Error("SQADD is in the generated table, the layering assumption broke")
	}
	if _, ok := LookupExtension(arch.ARM64, "SQADD"); !ok {
		t.Error("SQADD is missing from the extension layer")
	}
	if _, found := arch.ForArch(arch.ARM64).Lookup("ADD"); !found {
		t.Error("ADD vanished from the generated table")
	}
	if add, ok := LookupExtension(arch.ARM64, "ADD"); !ok || len(add) != 3 {
		t.Errorf("ADD carries %d extension forms, want 3", len(add))
	}
}

// TestEncodeExtensionGolden encodes through the registry and pins the same
// golden words the arch table tests pin, proving the registry resolves to the
// right encoding.
func TestEncodeExtensionGolden(t *testing.T) {
	for _, tt := range []struct {
		name string
		mnem string
		ops  []arch.ExtOperand
		want uint32
	}{
		{"unpredicated add", "ADD",
			[]arch.ExtOperand{
				arch.ExtVector(2, arch.ExtArrB), arch.ExtVector(0, arch.ExtArrB), arch.ExtVector(0, arch.ExtArrB),
			},
			0x04200040},
		{"predicated mul", "MUL",
			[]arch.ExtOperand{
				arch.ExtVector(0, arch.ExtArrB), arch.ExtPredicate(2, arch.ExtQualMerging), arch.ExtVector(0, arch.ExtArrB),
			},
			0x04100800},
		{"immediate add with derived shift", "ADD",
			[]arch.ExtOperand{arch.ExtImmediate(32512), arch.ExtVector(0, arch.ExtArrH)},
			0x2560efe0},
		{"signed immediate mul", "MUL",
			[]arch.ExtOperand{arch.ExtImmediate(-1), arch.ExtVector(0, arch.ExtArrB)},
			0x2530dfe0},
	} {
		got, err := EncodeExtension(arch.ARM64, tt.mnem, tt.ops...)
		if err != nil {
			t.Errorf("%s: encode: %v", tt.name, err)
			continue
		}
		if want := hex.EncodeToString([]byte{
			byte(tt.want), byte(tt.want >> 8), byte(tt.want >> 16), byte(tt.want >> 24),
		}); hex.EncodeToString(got) != want {
			t.Errorf("%s:\n  got  %x\n  want %s", tt.name, got, want)
		}
	}
}

// TestEncodeExtensionErrors checks the registry's diagnostics: a wrong arity
// names every form's count, an operand the first candidate rejects surfaces
// its own message once a later form takes over.
func TestEncodeExtensionErrors(t *testing.T) {
	if _, err := EncodeExtension(arch.ARM64, "ADD", arch.ExtVector(0, arch.ExtArrB)); err == nil {
		t.Error("one operand encoded, want an arity error")
	} else if !strings.Contains(err.Error(), "2 or 3 operands") {
		t.Errorf("arity error %q does not name the counts", err)
	}
	// The predicated candidate must answer for its own operands: the /Z
	// qualifier is rejected with the merging message, not the unpredicated
	// form's register-kind complaint.
	_, err := EncodeExtension(arch.ARM64, "ADD",
		arch.ExtVector(0, arch.ExtArrB), arch.ExtPredicate(0, arch.ExtQualZeroing), arch.ExtVector(0, arch.ExtArrB))
	if err == nil {
		t.Fatal("/Z encoded, want an error")
	}
	if !strings.Contains(err.Error(), "/M") {
		t.Errorf("error %q does not name the merging qualifier", err)
	}
	if _, err := EncodeExtension(arch.ARM64, "NOSUCHINSTR", arch.ExtVector(0, arch.ExtArrB)); err == nil ||
		!strings.Contains(err.Error(), "registers no extended instruction") {
		t.Errorf("unknown mnemonic error = %v", err)
	}
}

// TestExtensionEncodable checks the predicate the later Encodable hook will
// call: true exactly when the registry encodes the operand list.
func TestExtensionEncodable(t *testing.T) {
	if !ExtensionEncodable(arch.ARM64, "ADD",
		arch.ExtVector(0, arch.ExtArrS), arch.ExtVector(1, arch.ExtArrS), arch.ExtVector(2, arch.ExtArrS)) {
		t.Error("an encodable unpredicated add reported false")
	}
	if !ExtensionEncodable(arch.ARM64, "ADD",
		arch.ExtVector(1, arch.ExtArrS), arch.ExtPredicate(0, arch.ExtQualMerging), arch.ExtVector(0, arch.ExtArrS)) {
		t.Error("an encodable predicated add reported false")
	}
	if ExtensionEncodable(arch.ARM64, "ADD",
		arch.ExtVector(0, arch.ExtArrB), arch.ExtVector(0, arch.ExtArrS), arch.ExtVector(0, arch.ExtArrB)) {
		t.Error("mismatched arrangements reported encodable")
	}
	if ExtensionEncodable(arch.ARM64, "ADD", arch.ExtVector(0, arch.ExtArrB)) {
		t.Error("a one-operand add reported encodable")
	}
	if ExtensionEncodable(arch.ARM64, "NOSUCHINSTR") {
		t.Error("an unregistered mnemonic reported encodable")
	}
}

// TestExtensionArchIsolation is the architecture-binding negative case: the
// extension layer is registered for arm64 alone, and no other architecture
// answers its queries, not even for a mnemonic the amd64 base table carries.
func TestExtensionArchIsolation(t *testing.T) {
	ops := []arch.ExtOperand{
		arch.ExtVector(0, arch.ExtArrB), arch.ExtVector(0, arch.ExtArrB), arch.ExtVector(0, arch.ExtArrB),
	}
	for _, a := range []arch.Arch{arch.AMD64, arch.RISCV, arch.LOONG64, arch.Unknown} {
		if cands, ok := LookupExtension(a, "ADD"); ok || cands != nil {
			t.Errorf("LookupExtension(%s, ADD) offered %d candidates", a, len(cands))
		}
		if cands, ok := LookupExtension(a, "MUL"); ok || cands != nil {
			t.Errorf("LookupExtension(%s, MUL) offered %d candidates", a, len(cands))
		}
		if got, err := EncodeExtension(a, "ADD", ops...); err == nil {
			t.Errorf("EncodeExtension(%s, ADD) encoded %x, want a refusal", a, got)
		} else if !strings.Contains(err.Error(), string(a)) {
			t.Errorf("EncodeExtension(%s) error %q does not name the architecture", a, err)
		}
		if ExtensionEncodable(a, "ADD", ops...) {
			t.Errorf("ExtensionEncodable(%s, ADD) reported true", a)
		}
		if names := ExtensionNames(a); len(names) != 0 {
			t.Errorf("ExtensionNames(%s) = %v, want none", a, names)
		}
		if got := arch.Extensions(a); len(got) != 0 {
			t.Errorf("arch.Extensions(%s) carries %d instructions", a, len(got))
		}
	}
}

// TestExtensionNamesARM64 checks the completion-facing name list: every
// distinct mnemonic of the family, first-occurrence order, no duplicates.
func TestExtensionNamesARM64(t *testing.T) {
	want := []string{"ADD", "SUB", "SQADD", "UQADD", "SQSUB", "UQSUB", "MUL", "SMULH", "UMULH", "SUBR"}
	got := ExtensionNames(arch.ARM64)
	if strings.Join(got, ",") != strings.Join(want, ",") {
		t.Errorf("ExtensionNames(ARM64) = %v, want %v", got, want)
	}
	if n := len(arch.Extensions(arch.ARM64)); n != 23 {
		t.Errorf("the family registers %d instructions, want 23", n)
	}
}