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

package asm

import (
	"encoding/binary"
	"fmt"
	"os"
	"path/filepath"
	"slices"
	"strings"
	"testing"

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

// TestARM64SysRegsDifferential proves the whole system-register table against
// the toolchain at once: one TEXT whose body reads every register the table
// carries (and writes every writable one), assembled by gasm and by
// go tool asm, must agree byte for byte.  A single wrong op0/op1/CRn/CRm/op2
// packing names its register through the first differing word.
func TestARM64SysRegsDifferential(t *testing.T) {
	names := make([]string, 0, len(a64SysRegs))
	for name := range a64SysRegs {
		names = append(names, name)
	}
	slices.Sort(names)

	var body strings.Builder
	for i, name := range names {
		// R18 is the arm64 platform register and R29-R31 carry dedicated
		// meanings; a plain read/write destination keeps to R0-R17.
		reg := fmt.Sprintf("R%d", i%18)
		if a64SysRegs[name].read {
			body.WriteString(fmt.Sprintf("\tMRS %s, %s\n", name, reg))
		}
		if a64SysRegs[name].write {
			body.WriteString(fmt.Sprintf("\tMSR %s, %s\n", reg, name))
		}
	}
	src := "#include \"textflag.h\"\n\nTEXT ·sysregs(SB), NOSPLIT, $0\n" + body.String() + "\tRET\n"

	dir := t.TempDir()
	path := filepath.Join(dir, "sysregs_arm64.s")
	if err := os.WriteFile(path, []byte(src), 0o644); err != nil {
		t.Fatal(err)
	}
	assertARM64Differential(t, path, src, "sysregs")
}

// TestARM64FamiliesDifferential pins the non-sysreg families the arm64
// campaign added: the LSE compare-and-swap pairs, the VMOVI immediate, the
// SIMD narrow/long shift pairs, the VLD2/VLD3/VLD4 and VST2/VST3/VST4
// structure accesses with their post-index and replicate forms, LDPSW, the
// pointer-authentication hint and the DC maintenance operation.  Every
// spelling is the toolchain's own, taken from its arm64 testdata, and the
// bytes must agree word for word.
func TestARM64FamiliesDifferential(t *testing.T) {
	src := `#include "textflag.h"

TEXT ·families(SB), NOSPLIT, $0
	CASPD (R2, R3), (R2), (R8, R9)
	CASPW (R6, R7), (R8), (R4, R5)
	VMOVI $82, V0.B16
	VMOVI $146, V22.B16
	VSSHLL $0, V1.B8, V2.H8
	VSSHLL $7, V1.B8, V2.H8
	VSSHLL2 $0, V1.B16, V2.H8
	VSHRN $7, V1.H8, V0.B8
	VSHRN2 $31, V1.D2, V0.S4
	VLD2 (R29), [V23.H8, V24.H8]
	VLD2.P 16(R0), [V18.B8, V19.B8]
	VLD2.P (R1)(R2), [V15.S2, V16.S2]
	VLD3 (R27), [V11.S4, V12.S4, V13.S4]
	VLD3.P 48(RSP), [V11.S4, V12.S4, V13.S4]
	VLD4 (R15), [V10.H4, V11.H4, V12.H4, V13.H4]
	VLD4.P 32(R24), [V31.B8, V0.B8, V1.B8, V2.B8]
	VLD1R (R1), [V9.B8]
	VLD1R.P (R0), [V0.B16]
	VLD1R.P 2(R1), [V2.H4]
	VLD2R (R15), [V15.H4, V16.H4]
	VLD2R.P 16(R0), [V0.D2, V1.D2]
	VLD4R (R0), [V0.B8, V1.B8, V2.B8, V3.B8]
	VLD4R.P 16(RSP), [V31.S4, V0.S4, V1.S4, V2.S4]
	VST2 [V22.H8, V23.H8], (R23)
	VST2.P [V14.H4, V15.H4], 16(R17)
	VST2.P [V14.H4, V15.H4], (R3)(R17)
	VST3 [V1.D2, V2.D2, V3.D2], (R11)
	VST3.P [V18.S4, V19.S4, V20.S4], 48(R25)
	VST4 [V22.D2, V23.D2, V24.D2, V25.D2], (R3)
	VST4.P [V14.D2, V15.D2, V16.D2, V17.D2], 64(R15)
	LDPSW (R0), (R1, R2)
	LDPSW 4(R0), (R1, R2)
	LDPSW -4(R0), (R1, R2)
	PACIASP
	DC IVAC, R1
	RET
`
	dir := t.TempDir()
	path := filepath.Join(dir, "families_arm64.s")
	if err := os.WriteFile(path, []byte(src), 0o644); err != nil {
		t.Fatal(err)
	}
	assertARM64Differential(t, path, src, "families")
}

// assertARM64Differential assembles the same source with gasm and with the
// toolchain for arm64 and requires the named function's code bytes to agree.
// The live oracle is a deliberate-run comparison, so -short skips it (the
// push pipeline's mode); the golden bytes of the individual encoders are
// pinned separately in every mode.
func assertARM64Differential(t *testing.T, path, src, fn string) {
	t.Helper()
	oracle := oracleFuncCode(t, toolAsmObject(t, path, "arm64"))

	// The oracle keys its functions by the qualified object name
	// (pkg.name); match on the local part.
	want := map[string][]byte{}
	for name, code := range oracle {
		if _, after, ok := strings.Cut(name, "."); ok {
			want[after] = code
		} else {
			want[name] = code
		}
	}
	if want[fn] == nil {
		t.Fatalf("the oracle object carries no function %q (has %v)", fn, keysOf(want))
	}

	f, perrs := parser.Parse(path, src)
	if len(perrs) > 0 {
		t.Fatalf("parse: %v", perrs[0])
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("AssembleFileARM64: %v", err)
	}
	got := trimTrailingZeroWords(img.Code)
	wantB := trimTrailingZeroWords(want[fn])
	if len(got) != len(wantB) {
		t.Fatalf("gasm %d bytes, oracle %d bytes", len(got), len(wantB))
	}
	for i := range wantB {
		if got[i] != wantB[i] {
			t.Fatalf("word %d differs: gasm %08x, oracle %08x", i/4,
				binary.LittleEndian.Uint32(got[i:i+4]), binary.LittleEndian.Uint32(wantB[i:i+4]))
		}
	}
}

// trimTrailingZeroWords drops whole zero words off the end of a code span:
// an object pads a function to its alignment, and the raw image does not.
// A difference in the middle survives the trim untouched.
func trimTrailingZeroWords(b []byte) []byte {
	for len(b) >= 4 {
		last := b[len(b)-4:]
		if last[0]|last[1]|last[2]|last[3] != 0 {
			break
		}
		b = b[:len(b)-4]
	}
	return b
}