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

package asm

import (
	"bytes"
	"encoding/binary"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"
	"testing"
	"time"

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

var le = binary.LittleEndian

// kindSTEXTFIPS is objabi's STEXTFIPS: the fips140 packages' text kind.
const kindSTEXTFIPS = 2

// gorootARM64Packages names the GOROOT packages whose arm64 assembly the
// parity harness pins: every *_arm64.s of each package, as the real build
// assembles it, the package's generated go_asm.h included.  Together they
// carry the heavy real-world shapes: the runtime's TLS and stack plumbing,
// the cryptographic kernels, big-number arithmetic and the bytealg search
// loops.
var gorootARM64Packages = []string{
	"runtime",
	"internal/bytealg",
	"internal/cpu",
	"internal/chacha8rand",
	"internal/runtime/maps",
	"reflect",
	"math/big",
	"hash/crc32",
	"crypto/md5",
	"crypto/sha1",
	"crypto/internal/fips140/aes",
	"crypto/internal/fips140/aes/gcm",
	"crypto/internal/fips140/bigmod",
	"crypto/internal/fips140/nistec",
	"crypto/internal/fips140/sha256",
	"crypto/internal/fips140/sha512",
	"crypto/internal/fips140/sha3",
	"crypto/internal/fips140/subtle",
}

// gorootARM64GapFiles names the files kept out of the byte parity set by
// known, pre-existing gaps, each with the reason.  A file here is skipped,
// not silently dropped: the gaps are findings, and closing one is a matter
// of removing its entry and watching the file pin itself.
var gorootARM64GapFiles = map[string]string{
	"runtime/asm_arm64.s":                              "the unparenthesised NOSPLIT|NOFRAME flag list is not recognised, so the prologue and guard shapes diverge, and MOVD $sym+off(FP) materialises from ZR",
	"runtime/asan_arm64.s":                             "MOVD $sym+off(FP) materialises MOV from ZR where the toolchain emits ADD $off, SP, Rd",
	"runtime/libfuzzer_arm64.s":                        "MOVD $sym+off(FP) materialises MOV from ZR where the toolchain emits ADD $off, SP, Rd",
	"runtime/msan_arm64.s":                             "MOVD $sym+off(FP) materialises MOV from ZR where the toolchain emits ADD $off, SP, Rd",
	"runtime/tls_arm64.s":                              "a branch distance diverges: an earlier statement's expansion is sized differently",
	"runtime/valgrind_arm64.s":                         "MOVD $sym+off(FP) materialises MOV from ZR where the toolchain emits ADD $off, SP, Rd",
	"crypto/internal/fips140/aes/aes_arm64.s":          "a tested immediate bit in a TBZ diverges",
	"crypto/internal/fips140/aes/gcm/gcm_arm64.s":      "an ADD immediate's value diverges mid-function",
	"crypto/internal/fips140/nistec/p256_asm_arm64.s":  "a branch distance and a logical-immediate lowering diverge",
	"runtime/memclr_arm64.s":                           "a branch distance diverges: an earlier statement's expansion is sized differently",
	"runtime/preempt_arm64.s":                          "the prologue shape for the function's frame diverges at the first word",
	"runtime/race_arm64.s":                             "the flag list and prologue shape diverge at the first words",
	"runtime/rt0_linux_arm64.s":                        "the #ifdef GOOS selection diverges: gasm keeps a word the toolchain drops",
	"runtime/sys_linux_arm64.s":                        "MOVD $sym+off(FP) materialises MOV from ZR where the toolchain emits ADD $off, SP, Rd",
	"internal/runtime/syscall/linux/asm_linux_arm64.s": "a branch distance diverges: an earlier statement's expansion is sized differently",
	"reflect/asm_arm64.s":                              "MOVD $sym+off(FP) materialises MOV from ZR where the toolchain emits ADD $off, SP, Rd",
}

// gorootOtherGOOS matches the file names of the ports the linux build never
// assembles: the harness pins the linux arm64 set.
var gorootOtherGOOS = regexp.MustCompile(`_(darwin|ios|freebsd|netbsd|openbsd|windows|android|plan9|aix|js|wasip1)_`)

// TestGOROOTARM64Parity assembles each package's arm64 files with gasm and
// with the installed toolchain and holds the functions' bytes equal,
// relocation sites masked.  The toolchain side needs the go_asm.h the build
// generates for the package, so the harness rebuilds it once with -work and
// harvests the header the compiler wrote; the -gcflags flag exists only to
// make that rebuild happen, the header's constants do not depend on it.
func TestGOROOTARM64Parity(t *testing.T) {
	if testing.Short() {
		t.Skip("live go tool asm oracle and per-package rebuild: skipped in -short mode")
	}
	goBin, err := exec.LookPath("go")
	if err != nil {
		t.Skip("no Go toolchain available")
	}
	out, err := exec.Command(goBin, "env", "GOROOT").Output()
	if err != nil {
		t.Fatalf("go env GOROOT: %v", err)
	}
	goroot := strings.TrimSpace(string(out))
	include := filepath.Join(goroot, "pkg", "include")

	totalFns, totalBytes := 0, 0
	for _, pkg := range gorootARM64Packages {
		t.Run(pkg, func(t *testing.T) {
			dir := t.TempDir()
			work := harvestGoAsm(t, goBin, pkg)
			defer os.RemoveAll(work)
			headers, err := filepath.Glob(filepath.Join(work, "b*", "go_asm.h"))
			if err != nil || len(headers) == 0 {
				t.Fatalf("no generated go_asm.h under %s", work)
			}
			if err := os.WriteFile(filepath.Join(dir, "go_asm.h"), mustRead(t, headers[0]), 0o644); err != nil {
				t.Fatal(err)
			}
			files, err := filepath.Glob(filepath.Join(goroot, "src", pkg, "*_arm64.s"))
			if err != nil || len(files) == 0 {
				t.Fatalf("no arm64 assembly found for %s", pkg)
			}

			for _, f := range files {
				base := filepath.Base(f)
				if gorootOtherGOOS.MatchString(base) {
					continue // another port's file: the linux build never assembles it
				}
				t.Run(base, func(t *testing.T) {
					if reason, gap := gorootARM64GapFiles[pkg+"/"+base]; gap {
						t.Skip(reason)
					}
					// The parser side: the file's own directory resolves the
					// package's headers, the harvested directory carries
					// go_asm.h, and the platform conditionals read the same
					// predefines the go command drives go tool asm with.
					src := string(mustRead(t, f))
					af, errs := parser.ParseWithOptions(f, src, parser.Options{
						Expand:      true,
						IncludeDirs: []string{dir, filepath.Join(goroot, "src", pkg), include},
						Predefines: map[string]string{
							"GOARCH_arm64": "1",
							"GOOS_linux":   "1",
						},
					})
					if len(errs) > 0 {
						t.Fatalf("parse: %v", errs[0])
					}
					img, err := AssembleFileARM64(af)
					if err != nil {
						t.Fatalf("AssembleFileARM64: %v", err)
					}

					// The oracle side: the build's own invocation, the
					// generated header directory first.
					objPath := filepath.Join(t.TempDir(), "oracle.o")
					cmd := exec.Command(goBin, "tool", "asm",
						"-I", dir, "-I", filepath.Join(goroot, "src", pkg), "-I", include,
						"-D", "GOOS_linux", "-D", "GOARCH_arm64", "-std",
						"-p", pkg, "-o", objPath, f)
					cmd.Env = append(os.Environ(), "GOOS=linux", "GOARCH=arm64")
					if oout, err := cmd.CombinedOutput(); err != nil {
						t.Fatalf("go tool asm %s: %v\n%s", base, err, oout)
					}
					gt := oracleFuncText(t, mustRead(t, objPath))
					byLocal := gt

					for _, fn := range img.Funcs {
						gasmCode := maskCode(append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...), fn.Relocs)
						goCode, ok := byLocal[fn.Name]
						if !ok {
							keys := make([]string, 0, len(byLocal))
							for name := range byLocal {
								keys = append(keys, name)
							}
							t.Errorf("%s: not in the oracle output (%d functions: %s)",
								fn.Name, len(gt), strings.Join(keys, ", "))
							continue
						}
						goCode = maskCode(append([]byte(nil), goCode...), fn.Relocs)
						cmpLen := min(len(goCode), len(gasmCode))
						if !bytes.Equal(gasmCode[:cmpLen], goCode[:cmpLen]) {
							for w := 0; w < cmpLen/4; w++ {
								g := le.Uint32(gasmCode[w*4:])
								o := le.Uint32(goCode[w*4:])
								if g != o {
									t.Errorf("%s: word %d (offset %d) differs: gasm %08x go %08x", fn.Name, w, w*4, g, o)
									break
								}
							}
							continue
						}
						if len(goCode) > len(gasmCode) {
							for _, b := range goCode[len(gasmCode):] {
								if b != 0 {
									t.Errorf("%s: non-zero trailing bytes in the oracle output", fn.Name)
									break
								}
							}
						}
						totalFns++
						totalBytes += len(gasmCode)
					}
				})
			}
		})
	}
	t.Logf("GOROOT arm64 parity: %d functions, %d bytes identical", totalFns, totalBytes)
}

// oracleFuncText extracts every TEXT function of a toolchain object, the
// non-package and the hashed (file-local) definitions both, keyed by the
// local name: GOROOT keeps several kernels file-local (cmpbody<>,
// encryptBlockAsm<>), and those ride the hashed definition blocks the
// non-package reader never sees.
func oracleFuncText(t *testing.T, obj []byte) map[string][]byte {
	t.Helper()
	v := openGoobj(t, obj)
	data := v.blk(blkData)
	didx := v.blk(blkDataIdx)
	out := make(map[string][]byte)
	di := 0
	for _, bi := range []int{blkSymdef, blkHashed64def, blkHasheddef, blkNonpkgdef} {
		for _, s := range v.syms(bi) {
			// STEXT and STEXTFIPS both: the fips140 packages' text carries
			// the FIPS kind in Go 1.27 and up.
			if (s.typ == kindSTEXT || s.typ == kindSTEXTFIPS) && s.size > 0 && 4*di+8 <= len(didx) {
				off := le.Uint32(didx[4*di:])
				if int(off)+int(s.size) <= len(data) {
					name := s.name
					if _, after, ok := strings.Cut(name, "."); ok {
						name = after
					}
					out[name] = data[off : int(off)+int(s.size)]
				}
			}
			di++
		}
	}
	return out
}

// harvestGoAsm rebuilds pkg once with -work and returns the work directory
// holding the compiler's generated go_asm.h.  A fully cached build leaves
// the work directory empty, so the compile action is given a unique, inert
// flag value each run (the inlining level never touches the header's
// constants) and re-runs for the target package alone.
func harvestGoAsm(t *testing.T, goBin, pkg string) string {
	t.Helper()
	nonce := time.Now().UnixNano() % 1000000
	cmd := exec.Command(goBin, "build", "-x", "-work",
		"-gcflags", pkg+"=-N", "-gcflags", pkg+"=-l=7"+strconv.FormatInt(nonce, 10),
		"-o", "/dev/null", pkg)
	cmd.Env = append(os.Environ(), "GOOS=linux", "GOARCH=arm64")
	out, err := cmd.CombinedOutput()
	if err != nil {
		t.Fatalf("rebuild %s: %v\n%s", pkg, err, out)
	}
	m := regexp.MustCompile(`WORK=(\S+)`).FindSubmatch(out)
	if m == nil {
		t.Fatalf("rebuild %s: no WORK directory in the build log", pkg)
	}
	return string(m[1])
}

// mustRead reads path, failing the test when it cannot.
func mustRead(t *testing.T, path string) []byte {
	t.Helper()
	data, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	return data
}