// 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.
}

// TestArm64TLSLoad pins the local-exec TLS load: a symbol the file's own
// GLOBL marks TLSBSS loads as one MOVZ word carrying the R_ARM64_TLS_LE
// relocation (asm7.go case 69), the shape `go tool asm` emits for the
// runtime's tls_g accesses.  A non-TLS GLOBL keeps the ADRP+LDR pair.
func TestArm64TLSLoad(t *testing.T) {
	img := parseArm64File(t, "#include \"textflag.h\"\n\n"+
		"TEXT \u00b7f(SB), NOSPLIT, $0-0\n"+
		"\tMOVD tlsvar(SB), R0\n"+
		"\tMOVD plain(SB), R1\n"+
		"\tRET\n"+
		"GLOBL tlsvar(SB), TLSBSS, $8\n"+
		"GLOBL plain(SB), NOPTR, $8\n")
	fn := img.Funcs[0]
	if fn.Size != 4+8+4 {
		t.Fatalf("function size = %d, want 16", fn.Size)
	}
	if w := binary.LittleEndian.Uint32(img.Code[fn.Offset:]); w != 0xd2800000 {
		t.Errorf("TLS load word = %08x, want MOVZ 0 (d2800000)", w)
	}
	var tlsSeen, plainSeen bool
	for _, r := range fn.Relocs {
		if r.Name != "tlsvar" {
			continue
		}
		tlsSeen = true
		if r.Kind != RelArm64TLSLE {
			t.Errorf("tlsvar reloc kind = %v, want RelArm64TLSLE", r.Kind)
		}
		if r.Off != 0 {
			t.Errorf("tlsvar reloc off = %d, want 0", r.Off)
		}
	}
	for _, r := range fn.Relocs {
		if r.Name == "plain" && r.Kind == RelArm64LDST64 {
			plainSeen = true
		}
	}
	if !tlsSeen {
		t.Error("no tlsvar relocation recorded")
	}
	if !plainSeen {
		t.Error("the plain GLOBL load lost its ADRP+LDR relocation")
	}
	// The other widths have no TLS row: the toolchain refuses them.
	f, errs := parser.Parse("k_arm64.s", "TEXT \u00b7f(SB), NOSPLIT, $0-0\n"+
		"\tMOVW tlsvar(SB), R0\n"+
		"\tRET\n"+
		"GLOBL tlsvar(SB), TLSBSS, $8\n")
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	if _, err := AssembleFileARM64(f); err == nil {
		t.Error("MOVW of a TLS symbol assembled, want an illegal combination")
	}
}