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

package asm

import (
	"encoding/binary"
	"testing"

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

// parseArm64File is a helper assembling one arm64 source file.
func parseArm64File(t *testing.T, src string) *Image {
	t.Helper()
	f, errs := parser.Parse("k_arm64.s", src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("assemble: %v", err)
	}
	return img
}

// TestArm64RelocOffsetsIncludePrologue pins the function-relative relocation
// offsets of a framed function: the offsets used to exclude the prologue, so
// every relocation landed on a prologue instruction in the GOOBJ/ELF output.
// The function calls an external, so it is a non-leaf and carries the
// stack-split guard (12 bytes, small class) before the prologue.
func TestArm64RelocOffsetsIncludePrologue(t *testing.T) {
	img := parseArm64File(t, "TEXT \u00b7f(SB), $16-0\n"+
		"\tBL ext\u00b7foo(SB)\n"+
		"\tMOVD $gdata(SB), R5\n"+
		"\tMOVD $extsym(SB), R6\n"+
		"\tRET\n"+
		"GLOBL gdata(SB), $8\n")
	fn := img.Funcs[0]

	// Layout: 12-byte guard, 12-byte prologue, BL (24), ADRP+ADD (28, 32),
	// ADRP+ADD (36, 40), 12-byte epilogue with RET, 12-byte morestack block.
	want := []struct {
		off      int
		after    int
		name     string
		kind     RelocKind
		external bool
	}{
		{24, 28, "foo", RelArm64Branch, true},
		{28, 28, "gdata", RelArm64Addr, false},
		{32, 32, "gdata", RelArm64Addr, false},
		{36, 36, "extsym", RelArm64Addr, true},
		{40, 40, "extsym", RelArm64Addr, true},
		{60, 64, "runtime\u00b7morestack_noctxt", RelArm64Branch, true},
	}
	if len(fn.Relocs) != len(want) {
		t.Fatalf("relocs = %d, want %d", len(fn.Relocs), len(want))
	}
	for i, w := range want {
		r := fn.Relocs[i]
		if r.Off != w.off || r.After != w.after || r.Name != w.name || r.Kind != w.kind || r.External != w.external {
			t.Errorf("reloc %d = {off %d after %d name %q kind %d ext %v}, want {off %d after %d name %q kind %d ext %v}",
				i, r.Off, r.After, r.Name, r.Kind, r.External, w.off, w.after, w.name, w.kind, w.external)
		}
	}

	// The BL with a zero offset sits exactly at the first reloc site.
	code := img.Code[fn.Offset : fn.Offset+fn.Size]
	if w := binary.LittleEndian.Uint32(code[24:28]); w != 0x94000000 {
		t.Errorf("BL word = %08x, want 94000000", w)
	}
}

// TestArm64SBLoadStoreMatchesToolchain pins the ADRP scratch register
// (REGTMP, R27) and the LDST64 relocation kind for sym loads and stores,
// against the bytes go tool asm emits for MOVD sym(SB), R5.
func TestArm64SBLoadStoreMatchesToolchain(t *testing.T) {
	img := parseArm64File(t, "TEXT \u00b7ld(SB), NOSPLIT, $0\n"+
		"\tMOVD sym(SB), R5\n"+
		"\tMOVD R5, sym(SB)\n"+
		"\tRET\n"+
		"GLOBL sym(SB), $8\n")
	fn := img.Funcs[0]
	code := img.Code[fn.Offset : fn.Offset+fn.Size]

	// go tool asm: ADRP 0(PC), R27 (9000001b); MOVD (R27), R5 (f9400365);
	// ADRP 0(PC), R27; MOVD R5, (R27) (f9000365).
	for off, want := range map[int]uint32{0: 0x9000001b, 4: 0xf9400365, 8: 0x9000001b, 12: 0xf9000365} {
		if got := binary.LittleEndian.Uint32(code[off : off+4]); got != want {
			t.Errorf("word at %d = %08x, want %08x", off, got, want)
		}
	}

	if len(fn.Relocs) != 2 {
		t.Fatalf("relocs = %d, want 2", len(fn.Relocs))
	}
	for i, w := range []struct{ off, after int }{{0, 8}, {8, 16}} {
		r := fn.Relocs[i]
		if r.Kind != RelArm64LDST64 {
			t.Errorf("reloc %d kind = %d, want RelArm64LDST64 (%d)", i, r.Kind, RelArm64LDST64)
		}
		if r.Off != w.off || r.After != w.after {
			t.Errorf("reloc %d = {off %d after %d}, want {off %d after %d}", i, r.Off, r.After, w.off, w.after)
		}
	}
}

// TestArm64GOObjRelocTypes checks that GOOBJ emission succeeds with the new
// relocation kinds in play; the detailed layout is covered by the goobj tests.
func TestArm64GOObjRelocTypes(t *testing.T) {
	img := parseArm64File(t, "TEXT \u00b7ld(SB), NOSPLIT, $0\n"+
		"\tMOVD sym(SB), R5\n"+
		"\tMOVD R5, sym(SB)\n"+
		"\tRET\n"+
		"GLOBL sym(SB), $8\n")
	obj, err := img.GOObjectAARCH64("testpkg", "k_arm64.s")
	if err != nil {
		t.Fatalf("GOObjectAARCH64: %v", err)
	}
	if len(obj) == 0 {
		t.Fatal("empty object")
	}
	// The detailed layout is covered by the goobj tests; here we only pin
	// that emission succeeds with the new relocation kinds in play.
}