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

package main

import (
	"os"
	"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.
func generateFor(t *testing.T, dir, goarch string) string {
	t.Helper()
	hdrDir, err := generateGoAsmHeader(dir, goarch, t.TempDir())
	if err != nil {
		t.Fatalf("generateGoAsmHeader(%q, %s): %v", dir, 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)
	}
}

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)
	}
}

// TestGenerateGoAsmHeaderRuntime pins the generator against the real thing:
// the runtime package, whose header the toolchain's own -asmhdr output was
// sampled from.  Skipped in short mode: it type-checks the whole package.
func TestGenerateGoAsmHeaderRuntime(t *testing.T) {
	if testing.Short() {
		t.Skip("type-checks the whole runtime package")
	}
	dir, err := generateGoAsmHeader("/usr/local/go/src/runtime", "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)
		}
	}
}