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

package main

import (
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

// writePkg lays out a minimal Go package in a temp directory.
func writePkg(t *testing.T, files map[string]string) string {
	t.Helper()
	dir := t.TempDir()
	for name, src := range files {
		if err := os.WriteFile(filepath.Join(dir, name), []byte(src), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	return dir
}

// generateFor generates the header for dir and returns its text.  An empty
// goos means the ambient one.
func generateFor(t *testing.T, dir, goos, goarch string) string {
	t.Helper()
	hdrDir, err := generateGoAsmHeader(dir, goos, goarch, t.TempDir())
	if err != nil {
		t.Fatalf("generateGoAsmHeader(%q, %s, %s): %v", dir, goos, goarch, err)
	}
	b, err := os.ReadFile(filepath.Join(hdrDir, "go_asm.h"))
	if err != nil {
		t.Fatal(err)
	}
	return string(b)
}

func TestGenerateGoAsmHeaderShape(t *testing.T) {
	dir := writePkg(t, map[string]string{"sample.go": `package sample

const bufSize = 1024

const (
	a = iota * 8
	b
	c
)

const (
	strConst = "hello"
	boolConst = true
	floatConst = 1.5
	_ = "the blank identifier is skipped"
)

const shift = 1 << 20

type reader struct {
	r    int64
	w    int64
	_    [4]byte
	name string
}

type scalar int

type aliased struct {
	k uint32
	v uint32
}

type alias = aliased
`})
	hdr := generateFor(t, dir, "", "amd64")
	want := []string{
		"#define const_bufSize 1024",
		// iota resolves through go/types, one define per name.
		"#define const_a 0",
		"#define const_b 8",
		"#define const_c 16",
		`#define const_strConst "hello"`,
		"#define const_boolConst true",
		// Floats are the toolchain's own skip, as are blank names.
		"#define const_shift 1048576",
		// The blank field still occupies its bytes: the pad after w runs to
		// the string's 8-byte alignment.
		"#define reader__size 40",
		"#define reader_r 0",
		"#define reader_w 8",
		"#define reader_name 24",
		// Non-struct named types carry no defines; aliases to structs do.
		"#define aliased__size 8",
		"#define aliased_k 0",
		"#define aliased_v 4",
		"#define alias__size 8",
		"#define alias_k 0",
		"#define alias_v 4",
	}
	for _, w := range want {
		if !strings.Contains(hdr, w+"\n") {
			t.Errorf("header misses %q\ngot:\n%s", w, hdr)
		}
	}
	for _, banned := range []string{"#define const_floatConst", "#define _ ", "#define scalar"} {
		if strings.Contains(hdr, banned) {
			t.Errorf("header must not carry %s\ngot:\n%s", banned, hdr)
		}
	}
}

func TestGenerateGoAsmHeaderPerArch(t *testing.T) {
	dir := writePkg(t, map[string]string{
		"common.go": `package perarch

type layout struct {
	a int32
	p uintptr
}
`,
		// The build-tagged file set is part of the contract: a per-arch
		// package is exactly how internal/cpu declares its layouts.
		"const_amd64.go": `//go:build amd64

package perarch

const flavour = 1
`,
		"const_arm64.go": `//go:build arm64

package perarch

const flavour = 2
`,
	})
	amd64 := generateFor(t, dir, "", "amd64")
	arm64 := generateFor(t, dir, "", "arm64")
	if !strings.Contains(amd64, "#define const_flavour 1\n") {
		t.Errorf("amd64 header misses const_flavour 1:\n%s", amd64)
	}
	if !strings.Contains(arm64, "#define const_flavour 2\n") {
		t.Errorf("arm64 header misses const_flavour 2:\n%s", arm64)
	}
	if strings.Contains(arm64, "#define const_flavour 1\n") {
		t.Errorf("arm64 header must not carry the amd64 file's value")
	}
	// SizesFor makes the layout the target's: uintptr is 4 bytes wide on
	// 386 and 8 on amd64, which must move p and grow the struct.
	if !strings.Contains(amd64, "#define layout__size 16\n") || !strings.Contains(amd64, "#define layout_p 8\n") {
		t.Errorf("amd64 layout wrong:\n%s", amd64)
	}
	w386 := generateFor(t, dir, "", "386")
	if !strings.Contains(w386, "#define layout__size 8\n") || !strings.Contains(w386, "#define layout_p 4\n") {
		t.Errorf("386 layout wrong:\n%s", w386)
	}
}

// TestGenerateGoAsmHeaderGOOS pins the GOOS half of the target: only the
// platform's own files type-check into the header, which is why
// sys_darwin_arm64.s cannot assemble against a linux-generated one.
func TestGenerateGoAsmHeaderGOOS(t *testing.T) {
	dir := writePkg(t, map[string]string{
		"common.go": `package goosaware

type shared struct {
	a int32
}
`,
		"plat_darwin.go": `//go:build darwin

package goosaware

type platform struct {
	trampoline_numer int64
}
`,
		"plat_windows.go": `//go:build windows

package goosaware

type platform struct {
	callbackArgs__size int32
}
`,
	})
	darwin := generateFor(t, dir, "darwin", "arm64")
	if !strings.Contains(darwin, "#define platform__size 8\n") || !strings.Contains(darwin, "#define platform_trampoline_numer 0\n") {
		t.Errorf("darwin header misses the darwin layout:\n%s", darwin)
	}
	if strings.Contains(darwin, "callbackArgs") {
		t.Errorf("darwin header must not carry the windows layout:\n%s", darwin)
	}
	windows := generateFor(t, dir, "windows", "arm64")
	if !strings.Contains(windows, "#define platform_callbackArgs__size 0\n") {
		t.Errorf("windows header misses the windows layout:\n%s", windows)
	}
	if strings.Contains(windows, "trampoline_numer") {
		t.Errorf("windows header must not carry the darwin layout:\n%s", windows)
	}
	// The ambient GOOS is neither of the two, so only shared's defines are
	// emitted; the shared type keeps its layout there.
	ambient := generateFor(t, dir, "", "arm64")
	if !strings.Contains(ambient, "#define shared__size 4\n") {
		t.Errorf("ambient header misses the shared layout:\n%s", ambient)
	}
	if strings.Contains(ambient, "#define platform_") {
		t.Errorf("ambient header must not carry either platform layout:\n%s", ambient)
	}
}

func TestGoosFromFilename(t *testing.T) {
	for path, want := range map[string]string{
		"/x/sys_darwin_arm64.s":     "darwin",
		"/x/sys_windows_arm64.s":    "windows",
		"/x/asm_linux_amd64.s":      "linux",
		"/x/rt0_darwin_arm64.s":     "darwin",
		"/x/vgetrandom_zos_s390x.s": "zos",
		"/x/rt0_js_wasm.s":          "js",
		"/x/memmove_amd64.s":        "",
		"/x/vlop_arm.s":             "",
		"/x/stubs.s":                "",
	} {
		if got := goosFromFilename(path); got != want {
			t.Errorf("goosFromFilename(%q) = %q, want %q", path, got, want)
		}
	}
}

func TestGenerateGoAsmHeaderErrors(t *testing.T) {
	t.Run("type error", func(t *testing.T) {
		dir := writePkg(t, map[string]string{"bad.go": `package bad

const x = undefinedIdent
`})
		_, err := generateGoAsmHeader(dir, "", "amd64", t.TempDir())
		if err == nil {
			t.Fatal("generation must fail for a package that does not type-check")
		}
		if !strings.Contains(err.Error(), dir) {
			t.Errorf("error must name the package directory: %v", err)
		}
		if !strings.Contains(err.Error(), "type-check") {
			t.Errorf("error must say the package does not type-check: %v", err)
		}
	})
	t.Run("no go files", func(t *testing.T) {
		dir := t.TempDir()
		_, err := generateGoAsmHeader(dir, "", "amd64", t.TempDir())
		if err == nil {
			t.Fatal("generation must fail without Go files")
		}
		if !strings.Contains(err.Error(), dir) {
			t.Errorf("error must name the package directory: %v", err)
		}
	})
}

func TestNeedsGoAsmHeader(t *testing.T) {
	yes := "#include \"go_asm.h\"\n#include \"textflag.h\"\n"
	no := "#include \"textflag.h\"\n#include \"funcdata.h\"\n"
	if !needsGoAsmHeader(yes) {
		t.Error("needsGoAsmHeader(missing on a go_asm.h include)")
	}
	if needsGoAsmHeader(no) {
		t.Error("needsGoAsmHeader claims other headers need generation")
	}
}

func TestGoAsmHeaderResolved(t *testing.T) {
	dir := t.TempDir()
	if goAsmHeaderResolved(dir, nil) {
		t.Error("resolved with no header anywhere")
	}
	other := t.TempDir()
	if goAsmHeaderResolved(dir, []string{other}) {
		t.Error("resolved with an empty -I directory")
	}
	if err := os.WriteFile(filepath.Join(dir, "go_asm.h"), nil, 0o644); err != nil {
		t.Fatal(err)
	}
	if !goAsmHeaderResolved(dir, nil) {
		t.Error("not resolved with the header in the package directory")
	}
}

func TestOtherGOOSFile(t *testing.T) {
	for path, want := range map[string]bool{
		"/x/sys_windows_amd64.s": true,
		"/x/rt0_js_wasm.s":       true,
		"/x/sys_darwin_arm64.s":  true,
		"/x/sys_linux_amd64.s":   false,
		"/x/time_linux_amd64.s":  false,
		"/x/memmove_amd64.s":     false,
		"/x/generic.s":           false,
	} {
		if got := otherGOOSFile(path); got != want {
			t.Errorf("otherGOOSFile(%q) = %v, want %v", path, got, want)
		}
	}
}

// TestRunCorpusAuditGoAsm covers the audit wiring end to end: a package
// beside its kernel, the kernel living off the generated defines, and the
// histogram recording a generation failure as its own reason.
func TestRunCorpusAuditGoAsm(t *testing.T) {
	dir := t.TempDir()
	write := func(name, src string) {
		t.Helper()
		if err := os.WriteFile(filepath.Join(dir, name), []byte(src), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	write("pkg.go", `package corpus

const pageSize = 4096

type header struct {
	magic uint64
	flags uint64
}
`)
	write("kern_amd64.s", "#include \"go_asm.h\"\nTEXT \xc2\xb7f(SB), NOSPLIT, $0-16\n\tMOVQ\t$const_pageSize, AX\n\tMOVQ\t$header__size, BX\n\tRET\n")
	// The defines live in the file's own package; a kernel in a directory
	// without Go files has no package to generate from.
	if err := os.MkdirAll(filepath.Join(dir, "sub"), 0o755); err != nil {
		t.Fatal(err)
	}
	write(filepath.Join("sub", "lonely_arm64.s"), "#include \"go_asm.h\"\nTEXT \xc2\xb7g(SB), NOSPLIT, $0-0\n\tRET\n")

	stats, err := runCorpusAudit(dir, nil)
	if err != nil {
		t.Fatalf("runCorpusAudit: %v", err)
	}
	get := func(name string) *corpusTally {
		for i, tg := range stats.targets {
			if tg.name == name {
				return stats.tallies[i]
			}
		}
		t.Fatalf("no tally for %s", name)
		return nil
	}
	if a := get("amd64"); a.attempted != 1 || a.assembled != 1 {
		t.Errorf("amd64 = %d/%d, want 1/1", a.assembled, a.attempted)
	}
	// lonely_arm64.s is an arm64 file whose package cannot be generated.
	if a := get("arm64"); a.attempted != 1 || a.assembled != 0 {
		t.Errorf("arm64 = %d/%d, want 0/1", a.assembled, a.attempted)
	}
	if r := get("arm64").reasons["go_asm.h generation failed"]; r != 1 {
		t.Errorf("arm64 go_asm.h failure count = %d, want 1", r)
	}
}

// TestRunCorpusAuditGOOS covers the filename-derived GOOS end to end: a
// kernel whose name names darwin must have its header type-checked with
// GOOS=darwin, so the darwin-only constant it offsets with is defined.  The
// operand mirrors sys_darwin_arm64.s's trampoline, where a missing define
// leaves an unexpanded symbol in the offset and fails.
func TestRunCorpusAuditGOOS(t *testing.T) {
	dir := t.TempDir()
	write := func(name, src string) {
		t.Helper()
		if err := os.WriteFile(filepath.Join(dir, name), []byte(src), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	write("pkg.go", "package corpus\n")
	write("plat_darwin.go", "//go:build darwin\n\npackage corpus\n\nconst trampolineNumer = 8\n")
	write("kern_darwin_arm64.s", "#include \"go_asm.h\"\n"+
		"GLOBL timebase<>(SB), NOPTR, $16\n"+
		"TEXT \xc2\xb7g(SB), NOSPLIT, $0-0\n"+
		"\tMOVD\ttimebase<>+const_trampolineNumer(SB), R0\n"+
		"\tRET\n")

	stats, err := runCorpusAudit(dir, nil)
	if err != nil {
		t.Fatalf("runCorpusAudit: %v", err)
	}
	var arm *corpusTally
	for i, tg := range stats.targets {
		if tg.name == "arm64" {
			arm = stats.tallies[i]
		}
	}
	if arm == nil {
		t.Fatal("no arm64 tally")
	}
	if arm.attempted != 1 || arm.assembled != 1 {
		t.Errorf("arm64 = %d/%d, want 1/1; reasons: %v", arm.assembled, arm.attempted, arm.reasons)
	}
}

// TestRunCorpusAuditBuildConstraint covers the //go:build classification end
// to end: a generic-named file whose constraint admits one target is
// attempted there alone (cpu_x86.s on amd64), and a file whose constraint
// admits none of the four targets is never attempted (the msan and
// goexperiment trees).
func TestRunCorpusAuditBuildConstraint(t *testing.T) {
	dir := t.TempDir()
	write := func(name, src string) {
		t.Helper()
		if err := os.WriteFile(filepath.Join(dir, name), []byte(src), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	write("x86.s", "//go:build 386 || amd64\n\nTEXT \xc2\xb7f(SB), NOSPLIT, $0\n\tRET\n")
	write("racey.s", "//go:build race\n\nTEXT \xc2\xb7r(SB), NOSPLIT, $0\n\tRET\n")
	write("plain.s", "TEXT \xc2\xb7p(SB), NOSPLIT, $0\n\tRET\n")

	stats, err := runCorpusAudit(dir, nil)
	if err != nil {
		t.Fatalf("runCorpusAudit: %v", err)
	}
	tally := func(name string) *corpusTally {
		for i, tg := range stats.targets {
			if tg.name == name {
				return stats.tallies[i]
			}
		}
		t.Fatalf("no tally for %s", name)
		return nil
	}
	if stats.narrowed != 1 || stats.excluded != 1 || stats.generic != 1 {
		t.Errorf("buckets = narrowed %d, excluded %d, generic %d; want 1, 1, 1", stats.narrowed, stats.excluded, stats.generic)
	}
	if a := tally("amd64"); a.attempted != 2 || a.assembled != 2 {
		t.Errorf("amd64 = %d/%d, want 2/2 (x86.s and plain.s)", a.assembled, a.attempted)
	}
	for _, name := range []string{"arm64", "riscv64", "loong64"} {
		if a := tally(name); a.attempted != 1 || a.assembled != 1 {
			t.Errorf("%s = %d/%d, want 1/1 (plain.s only)", name, a.assembled, a.attempted)
		}
	}
	if stats.full != 2 {
		t.Errorf("full = %d, want 2 (x86.s over its one target, plain.s over all four)", stats.full)
	}
}

// TestGenerateGoAsmHeaderRuntime pins the generator against the real thing:
// the runtime package of the ambient toolchain, whose header the toolchain's
// own -asmhdr output was sampled from.  Skipped in short mode: it type-checks
// the whole package.  The GOROOT comes from the go command itself, so the
// test follows whatever toolchain the host provides.
func TestGenerateGoAsmHeaderRuntime(t *testing.T) {
	if testing.Short() {
		t.Skip("type-checks the whole runtime package")
	}
	out, err := exec.Command("go", "env", "GOROOT").Output()
	if err != nil {
		t.Skipf("no Go toolchain: %v", err)
	}
	runtimeDir := filepath.Join(strings.TrimSpace(string(out)), "src", "runtime")
	dir, err := generateGoAsmHeader(runtimeDir, "", "amd64", t.TempDir())
	if err != nil {
		t.Fatalf("generateGoAsmHeader(runtime): %v", err)
	}
	b, err := os.ReadFile(dir + "/go_asm.h")
	if err != nil {
		t.Fatal(err)
	}
	hdr := string(b)
	for _, want := range []string{
		"#define const_hashSize 8\n",
		"#define const_avxSupported 1\n",
		"#define const_pageSize 8192\n",
		"#define g_stackguard0 16\n",
		"#define m__size ",
	} {
		if !strings.Contains(hdr, want) {
			t.Errorf("runtime header misses %q", want)
		}
	}
}