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

package asm

import (
	"bytes"
	"encoding/hex"
	"testing"

	"sourcedock.dev/petrbalvin/gasm-devkit/ast"
	"sourcedock.dev/petrbalvin/gasm-devkit/parser"
)

// The expected bytes are pinned from `go tool asm` output (Go 1.27, amd64,
// verified with go tool objdump): the stack-split guard classes, the morestack
// block and the auto-NOSPLIT leaf behaviour.
func TestStackGuardBytes(t *testing.T) {
	for _, tt := range []struct {
		name string
		src  string
		want string
	}{
		{"leafsmall", "TEXT \u00b7leafsmall(SB), $16-0\n\tRET\n",
			"554889e54883ec104883c4105dc3"},
		{"leafmed", "TEXT \u00b7leafmed(SB), $256-0\n\tRET\n",
			"644c8b3425000000004c8da42478ffffff4d3b66107614554889e54881ec000100004881c4000100005dc3e800000000ebce"},
		{"leafbig", "TEXT \u00b7leafbig(SB), $8192-0\n\tRET\n",
			"644c8b3425000000004989e44981ec881f0000721a4d3b66107614554889e54881ec002000004881c4002000005dc3e800000000ebca"},
		{"callsmall", "TEXT \u00b7callsmall(SB), $16-0\n\tCALL \u00b7other(SB)\n\tRET\nTEXT \u00b7other(SB), NOSPLIT, $0\n\tRET\n",
			"644c8b342500000000493b66107613554889e54883ec10e8000000004883c4105dc3e800000000ebd7"},
		{"nosplit", "TEXT \u00b7nosplit(SB), NOSPLIT, $16-0\n\tRET\n",
			"554889e54883ec104883c4105dc3"},
	} {
		f, errs := parser.Parse("g_amd64.s", tt.src)
		if len(errs) > 0 {
			t.Fatalf("%s: parse: %v", tt.name, errs)
		}
		img, err := AssembleFile(f)
		if err != nil {
			t.Fatalf("%s: assemble: %v", tt.name, err)
		}
		fn := img.Funcs[0]
		// The toolchain's object leaves every relocation field zero for the
		// linker, while the gasm image resolves file-internal references, so
		// the comparison masks the patch sites the way verify's ground truth
		// does.
		code := append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...)
		for _, r := range fn.Relocs {
			for j := r.Off; j < r.Off+4 && j < len(code); j++ {
				code[j] = 0
			}
		}
		got := hex.EncodeToString(code)
		if got != tt.want {
			t.Errorf("%s:\n  got  %s\n  want %s", tt.name, got, tt.want)
		}
	}
}

// TestStackGuardRelocs checks the guard's patch sites: the TLS slot and the
// morestack call.
func TestStackGuardRelocs(t *testing.T) {
	f, errs := parser.Parse("g_amd64.s", "TEXT \u00b7f(SB), $256-0\n\tRET\n")
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("assemble: %v", err)
	}
	relocs := img.Funcs[0].Relocs
	if len(relocs) != 2 {
		t.Fatalf("relocs = %d, want 2", len(relocs))
	}
	tls, call := relocs[0], relocs[1]
	if tls.Kind != RelTLSLE || tls.Off != 5 || tls.Name != "" || tls.External {
		t.Errorf("tls reloc = %+v, want RelTLSLE at 5 with no symbol", tls)
	}
	if call.Kind != RelCall || call.Name != "runtime\u00b7morestack_noctxt" || !call.External {
		t.Errorf("call reloc = %+v, want RelCall to runtime.morestack_noctxt", call)
	}
}

// TestStackGuardGOObj emissions succeed with the guard's TLS and builtin
// references in play.
func TestStackGuardGOObj(t *testing.T) {
	f, errs := parser.Parse("g_amd64.s", "TEXT \u00b7f(SB), $256-0\n\tCALL \u00b7helper(SB)\n\tRET\nTEXT \u00b7helper(SB), NOSPLIT, $0\n\tRET\n")
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("assemble: %v", err)
	}
	obj, err := img.GOObject("testpkg", "g_amd64.s")
	if err != nil {
		t.Fatalf("GOObject: %v", err)
	}
	if !bytes.Contains(obj, []byte("go120ld")) {
		t.Fatal("object lacks the GOOBJ magic")
	}
}

// The arm64 stack-split guard, pinned from `go tool asm` (Go 1.27, arm64):
// the guard classes, the auto-NOSPLIT leaf behaviour and the morestack
// block.  Relocation fields are masked: the toolchain's object leaves them
// zero for the linker, the gasm image resolves file-internal references.
func TestStackGuardBytesARM64(t *testing.T) {
	for _, tt := range []struct {
		name string
		src  string
		want string
	}{
		{"leafsmall", "TEXT \u00b7leafsmall(SB), $16-0\n\tRET\n",
			"fe0f1ef8fd831ff8fd2300d1fd630091ff830091c0035fd6"},
		{"leafmed", "TEXT \u00b7leafmed(SB), $256-0\n\tRET\n",
			"900b40f9f14302d13f0210eb09010054f44304d19dfa3fa99f020091fd2300d1fd230491ff430491c0035fd6e3031eaa00000000f3ffff17"},
		{"leafbig", "TEXT \u00b7leafbig(SB), $8192-0\n\tRET\n",
			"900b40f91bf283d2f1633beba30100543f0210eb690100541b0284d2f4633bcb9dfa3fa99f020091fd2300d11b0184d2fd633b8b1b0284d2ff633b8bc0035fd6e3031eaa00000000eeffff17"},
		{"callsmall", "TEXT \u00b7callsmall(SB), $16-0\n\tCALL \u00b7other(SB)\n\tRET\nTEXT \u00b7other(SB), NOSPLIT, $0\n\tRET\n",
			"900b40f9ff6330eb09010054fe0f1ef8fd831ff8fd2300d100000000fd835ff8fe0742f8c0035fd6e3031eaa00000000f4ffff17"},
		{"nosplit", "TEXT \u00b7nosplit(SB), NOSPLIT, $16-0\n\tRET\n",
			"fe0f1ef8fd831ff8fd2300d1fd630091ff830091c0035fd6"},
	} {
		f, errs := parser.Parse("g_arm64.s", tt.src)
		if len(errs) > 0 {
			t.Fatalf("%s: parse: %v", tt.name, errs)
		}
		img, err := AssembleFileARM64(f)
		if err != nil {
			t.Fatalf("%s: assemble: %v", tt.name, err)
		}
		fn := img.Funcs[0]
		code := append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...)
		for _, r := range fn.Relocs {
			for j := r.Off; j < r.Off+4 && j < len(code); j++ {
				code[j] = 0
			}
		}
		got := hex.EncodeToString(code)
		if got != tt.want {
			t.Errorf("%s:\n  got  %s\n  want %s", tt.name, got, tt.want)
		}
	}
}

// The riscv64 stack-split guard, pinned from `go tool asm` (Go 1.27,
// riscv64): the morestack call sits between the guard and the body, and the
// guard branches forward over it.  Relocation fields are masked.
func TestStackGuardBytesRISCV64(t *testing.T) {
	for _, tt := range []struct {
		name string
		src  string
		want string
	}{
		{"leafsmall", "TEXT \u00b7leafsmall(SB), $16-0\n\tRET\n",
			"03b30d0163662300000000006ff05fff233411fe211106e08260610167800000"},
		{"leafmed", "TEXT \u00b7leafmed(SB), $256-0\n\tRET\n",
			"03b30d01930381f763667300000000006ff01fff233c11ee130181ef06e082601301811067800000"},
		{"leafbig", "TEXT \u00b7leafbig(SB), $8192-0\n\tRET\n",
			"03b30d0189639b8383f863697100f97f9b8f8f07b303f10163667300000000006ff01ffef97f8a9f23bc1ffef97fe13f7e9106e08260896fa12f7e9167800000"},
		{"frameless", "TEXT \u00b7frameless(SB), $0-0\n\tCALL \u00b7other(SB)\n\tRET\nTEXT \u00b7other(SB), NOSPLIT, $0\n\tRET\n",
			"03b30d0163662300000000006ff05fff233c11fe611106e0000000008260210167800000"},
		{"nosplit", "TEXT \u00b7nosplit(SB), NOSPLIT, $16-0\n\tRET\n",
			"233411fe211106e08260610167800000"},
	} {
		f, errs := parser.Parse("g_riscv64.s", tt.src)
		if len(errs) > 0 {
			t.Fatalf("%s: parse: %v", tt.name, errs)
		}
		img, err := AssembleFileRISCV(f)
		if err != nil {
			t.Fatalf("%s: assemble: %v", tt.name, err)
		}
		fn := img.Funcs[0]
		code := append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...)
		for _, r := range fn.Relocs {
			for j := r.Off; j < r.Off+4 && j < len(code); j++ {
				code[j] = 0
			}
		}
		got := hex.EncodeToString(code)
		if got != tt.want {
			t.Errorf("%s:\n  got  %s\n  want %s", tt.name, got, tt.want)
		}
	}
}

// The loong64 stack-split guard, pinned from `go tool asm` (Go 1.27,
// loong64): every guard class (including the medium class with the
// materialised constant and the big class with the ORI-less constants), the
// auto-NOSPLIT leaf behaviour, the large-frame R30 prologue/epilogue forms
// and the morestack block.  Relocation fields are masked.
func TestStackGuardBytesLOONG64(t *testing.T) {
	for _, tt := range []struct {
		name string
		src  string
		want string
	}{
		{"leafsmall", "TEXT \u00b7leafsmall(SB), $16-0\n\tRET\n",
			"61a0ff2963a0ff026100c0296360c0022000004c"},
		{"leafmed", "TEXT \u00b7leafmed(SB), $256-0\n\tRET\n",
			"d442c02878e0fd0294e21200801a004061e0fb2963e0fb026100c0296320c4022000004c3f00150000000000ffd7ff53"},
		{"nosplit", "TEXT \u00b7nosplit(SB), NOSPLIT, $16-0\n\tRET\n",
			"61a0ff2963a0ff026100c0296360c0022000004c"},
		// The LR store leaves the 12-bit store-offset range while the SP
		// adjust immediate still fits, and the epilogue adjusts through a
		// single ORI.
		{"fit2048", "TEXT \u00b7fit2048(SB), $2040-0\n\tRET\n",
			"d442c0287800e20294e21200802600401e000014de8f1000c103e0296300e0026100c0291e00a00363f810002000004c3f00150000000000ffcbff53"},
		// Medium class at the materialisation boundary (off = 2048 still
		// immediate, 2049+ goes through R30).
		{"med2048off", "TEXT \u00b7med2048off(SB), $2168-0\n\tRET\n",
			"d442c0287800e00294e21200802e0040feffff15de8f1000c103de29feffff15de039e0363f810006100c0291e00a20363f810002000004c3f00150000000000ffc3ff53"},
		{"medmat", "TEXT \u00b7medmat(SB), $2176-0\n\tRET\n",
			"d442c028feffff15dee39f0378f8100094e21200802e0040feffff15de8f1000c1e3dd29feffff15dee39d0363f810006100c0291e20a20363f810002000004c3f00150000000000ffbbff53"},
		// Big class with the rounding-split store and the floor-split adjust.
		{"leafbig", "TEXT \u00b7leafbig(SB), $8192-0\n\tRET\n",
			"d442c0283e000014de23be0378f8120000470044deffff15dee3810378f8100094e2120080320040deffff15de8f1000c1e3ff29beffff15dee3bf0363f810006100c0295e000014de23800363f810002000004c3f00150000000000ffa7ff53"},
		// Zero low 12 bits drop the ORI from the store, the adjust and the
		// epilogue materialisation.
		{"bigzero", "TEXT \u00b7bigzero(SB), $4088-0\n\tRET\n",
			"d442c028feffff15de03820378f8100094e21200802a0040feffff15de8f1000c103c029feffff1563f810006100c0293e00001463f810002000004c3f00150000000000ffbfff53"},
		// Big class whose first constant has a zero high part: a single ORI.
		{"big3976", "TEXT \u00b7big3976(SB), $4096-0\n\tRET\n",
			"d442c0281e20be0378f8120000470044feffff15dee3810378f8100094e2120080320040feffff15de8f1000c1e3ff29deffff15dee3bf0363f810006100c0293e000014de23800363f810002000004c3f00150000000000ffabff53"},
		// Big class at a multiple of 4096: both guard constants lose their
		// ORI word.
		{"giantlo0", "TEXT \u00b7giantlo0(SB), $4216-0\n\tRET\n",
			"d442c0283e00001478f8120000430044feffff1578f8100094e2120080320040feffff15de8f1000c103fe29deffff15de03be0363f810006100c0293e000014de03820363f810002000004c3f00150000000000ffafff53"},
		// Non-leaf big frame: the body call plus the LR restore epilogue.
		{"callbig", "TEXT \u00b7callbig(SB), $8192-0\n\tCALL \u00b7other(SB)\n\tRET\nTEXT \u00b7other(SB), NOSPLIT, $0\n\tRET\n",
			"d442c0283e000014de23be0378f81200004f0044deffff15dee3810378f8100094e21200803a0040deffff15de8f1000c1e3ff29beffff15dee3bf0363f810006100c029000000006100c0285e000014de23800363f810002000004c3f00150000000000ff9fff53"},
	} {
		f, errs := parser.Parse("g_loong64.s", tt.src)
		if len(errs) > 0 {
			t.Fatalf("%s: parse: %v", tt.name, errs)
		}
		img, err := AssembleFileLOONG64(f)
		if err != nil {
			t.Fatalf("%s: assemble: %v", tt.name, err)
		}
		fn := img.Funcs[0]
		code := append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...)
		for _, r := range fn.Relocs {
			for j := r.Off; j < r.Off+4 && j < len(code); j++ {
				code[j] = 0
			}
		}
		got := hex.EncodeToString(code)
		if got != tt.want {
			t.Errorf("%s:\n  got  %s\n  want %s", tt.name, got, tt.want)
		}
	}
}

// TestStackGuardGOObjInternalCall checks that GOOBJ emission succeeds when a
// guarded function calls a TEXT symbol of the same file, for every arch's
// call relocation kind.
func TestStackGuardGOObjInternalCall(t *testing.T) {
	for _, tt := range []struct {
		src      string
		assemble func(*ast.File) (*Image, error)
	}{
		{"g_amd64.s", AssembleFile},
		{"g_arm64.s", AssembleFileARM64},
		{"g_riscv64.s", AssembleFileRISCV},
		{"g_loong64.s", AssembleFileLOONG64},
	} {
		f, errs := parser.Parse(tt.src, "TEXT \u00b7callsmall(SB), $16-0\n\tCALL \u00b7other(SB)\n\tRET\nTEXT \u00b7other(SB), NOSPLIT, $0\n\tRET\n")
		if len(errs) > 0 {
			t.Fatalf("%s: parse: %v", tt.src, errs)
		}
		img, err := tt.assemble(f)
		if err != nil {
			t.Fatalf("%s: assemble: %v", tt.src, err)
		}
		if _, err := img.GOObject("testpkg", tt.src); err != nil {
			t.Errorf("%s: GOObject: %v", tt.src, err)
		}
	}
}