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

package verify

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"slices"
	"strings"
	"testing"

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

// TestGroundTruthARM64 assembles the arm64 test kernels with gasm and
// compares them byte-for-byte against `go tool asm` (GOARCH=arm64).  The
// relocation fields of static-symbol references are masked before the
// comparison, since the toolchain leaves them zero for the linker.
func TestGroundTruthARM64(t *testing.T) {
	for _, path := range []string{
		"../testdata/verify/basic_arm64.s",
		"../testdata/verify/fp_arm64.s",
		"../testdata/verify/movimm_arm64.s",
		"../testdata/verify/branch_arm64.s",
		"../testdata/verify/call_arm64.s",
		"../testdata/verify/bigframe_arm64.s",
		"../testdata/verify/guard_arm64.s",
		"../testdata/verify/indirect_arm64.s",
		"../testdata/verify/exclusive_arm64.s",
		"../testdata/verify/shifts_arm64.s",
		"../testdata/verify/atomics_arm64.s",
		"../testdata/verify/qmov_arm64.s",
		"../testdata/verify/splits_arm64.s",
		"../testdata/verify/regoffset_arm64.s",
		"../testdata/verify/simdarr_arm64.s",
		"../testdata/verify/crypto_arm64.s",
		"../testdata/verify/integer_arm64.s",
		"../testdata/verify/simd_arm64.s",
		"../testdata/verify/widenimm_arm64.s",
		"../testdata/verify/carryshift_arm64.s",
		"../testdata/verify/simdmove_arm64.s",
		"../testdata/verify/bookkeep_arm64.s",
		"../testdata/verify/system_arm64.s",
	} {
		t.Run(path, func(t *testing.T) {
			src, err := os.ReadFile(path)
			if err != nil {
				t.Fatalf("read: %v", err)
			}
			f, errs := parser.Parse(path, string(src))
			if len(errs) > 0 {
				t.Fatalf("parse: %v", errs)
			}
			img, err := asm.AssembleFileARM64(f)
			if err != nil {
				t.Fatalf("AssembleFileARM64: %v", err)
			}
			gt, err := GroundTruthARM64(path)
			if err != nil {
				t.Fatalf("GroundTruthARM64: %v", err)
			}

			matched := 0
			for _, fn := range img.Funcs {
				gasmCode := maskRelocs(append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...), fn.Relocs)
				goCode, ok := gt[fn.Name]
				if !ok {
					t.Errorf("%s: not in ground truth (%d functions)", fn.Name, len(gt))
					continue
				}
				goCode = maskRelocs(goCode, fn.Relocs)
				// The Go toolchain may add zero padding at the end of
				// functions.  Compare up to the shorter length, then
				// verify any trailing bytes are zero.
				cmpLen := min(len(goCode), len(gasmCode))
				if !bytes.Equal(gasmCode[:cmpLen], goCode[:cmpLen]) {
					t.Errorf("%s: MISMATCH gasm=%d go=%d bytes\n%s", fn.Name, len(gasmCode), len(goCode), diffHex(gasmCode, goCode))
					continue
				}
				// Check trailing padding is zero.
				trailingOK := true
				if len(goCode) > len(gasmCode) {
					for _, b := range goCode[len(gasmCode):] {
						if b != 0 {
							trailingOK = false
							break
						}
					}
				}
				if !trailingOK {
					t.Errorf("%s: non-zero trailing bytes in go tool asm output", fn.Name)
					continue
				}
				matched++
				t.Logf("%s: MATCH (%d bytes, go=%d)", fn.Name, fn.Size, len(goCode))
			}
			if matched == 0 {
				t.Fatal("no functions matched")
			}
		})
	}
}

// TestGroundTruthARM64OffsetSweep sweeps load/store offsets across every
// band boundary the encoder branches on and compares the assembled bytes
// against the toolchain for the whole sweep in one go tool asm run.  The
// offsets the toolchain itself rejects are dropped first, so the comparison
// pins the split and single-instruction encodings, not the error parity.
func TestGroundTruthARM64OffsetSweep(t *testing.T) {
	if testing.Short() {
		t.Skip("live go tool asm oracle: 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)
	}
	includeDir := filepath.Join(strings.TrimSpace(string(out)), "pkg", "include")

	bands := []int64{
		0, 255, 256, 4094, 4095, 4096, 4097, 8190, 8191, 8192,
		16380, 16384, 32760, 32768, 65504, 65520, 65536, 65552,
		0xfff000, 0xfff000 + 8190, 0xfff000 + 16380, 0xfff000 + 32760,
		0xfff000 + 65520, 0xfffff8, 0xffffff, 0x1000000, 0x1006ff8,
		0xfff000 + 0xfff<<4, 0x1000ffe,
	}
	var offs []int64
	for _, b := range bands {
		for _, d := range []int64{-8, -2, -1, 0, 1, 2, 8} {
			if v := b + d; v >= 0 {
				offs = append(offs, v)
			}
		}
		offs = append(offs, -b-1, -b, -b+1)
	}
	slices.Sort(offs)
	offs = slices.Compact(offs)

	// One function per mnemonic; a line the toolchain rejects (an unaligned
	// band edge, a pool case gasm rejects until the pool lands) would fail
	// the whole oracle run, so the sweep keeps only the offsets both
	// assemblers accept: build per-mnemonic files and compare the ones that
	// assemble on both sides.
	for _, mnem := range []string{"MOVB", "MOVH", "MOVW", "MOVD", "FMOVS", "FMOVD", "FMOVQ", "LDP", "STP"} {
		t.Run(mnem, func(t *testing.T) {
			var body strings.Builder
			for _, off := range offs {
				switch mnem {
				case "LDP":
					fmt.Fprintf(&body, "\t%s\t%d(R2), (R1, R3)\n", mnem, off)
				case "STP":
					fmt.Fprintf(&body, "\t%s\t(R1, R3), %d(R2)\n", mnem, off)
				case "FMOVS", "FMOVD", "FMOVQ":
					fmt.Fprintf(&body, "\t%s\tF1, %d(R2)\n", mnem, off)
				default:
					fmt.Fprintf(&body, "\t%s\tR1, %d(R2)\n", mnem, off)
				}
			}
			src := "TEXT ·f(SB), 7, $0-0\n" + body.String() + "\tRET\n"
			dir := t.TempDir()
			path := filepath.Join(dir, "sweep.s")
			if err := os.WriteFile(path, []byte(src), 0o644); err != nil {
				t.Fatal(err)
			}
			f, errs := parser.Parse(path, src)
			if len(errs) > 0 {
				t.Fatalf("parse: %v", errs)
			}
			gasmImg, gasmErr := asm.AssembleFileARM64(f)
			objPath := filepath.Join(dir, "oracle.o")
			cmd := exec.Command(goBin, "tool", "asm", "-I", includeDir, "-p", "sweep", "-o", objPath, path)
			cmd.Env = append(os.Environ(), "GOARCH=arm64")
			oracleOut, oracleErr := cmd.CombinedOutput()
			if gasmErr != nil && oracleErr != nil {
				return // both reject: agreement on the boundary
			}
			if gasmErr != nil {
				t.Fatalf("gasm rejects what the toolchain accepts:\n%v", gasmErr)
			}
			if oracleErr != nil {
				t.Fatalf("gasm accepts what the toolchain rejects:\n%s", oracleOut)
			}
			gt, err := GroundTruthARM64(path)
			if err != nil {
				t.Fatalf("ground truth: %v", err)
			}
			goCode, ok := gt["f"]
			if !ok {
				t.Fatal("oracle function missing")
			}
			for _, fn := range gasmImg.Funcs {
				if fn.Name != "f" {
					continue
				}
				gasmCode := gasmImg.Code[fn.Offset : fn.Offset+fn.Size]
				cmpLen := min(len(goCode), len(gasmCode))
				if !bytes.Equal(gasmCode[:cmpLen], goCode[:cmpLen]) {
					for w := 0; w < cmpLen/4; w++ {
						g := binary.LittleEndian.Uint32(gasmCode[w*4 : w*4+4])
						o := binary.LittleEndian.Uint32(goCode[w*4 : w*4+4])
						if g != o {
							t.Fatalf("word %d (offset %d) differs: gasm %#08x go %#08x", w, w*4, g, o)
						}
					}
					t.Fatalf("prefixes equal but lengths differ (gasm %d, go %d)", len(gasmCode), len(goCode))
				}
				for _, b := range goCode[len(gasmCode):] {
					if b != 0 {
						t.Fatalf("non-zero trailing oracle bytes (gasm %d, go %d)", len(gasmCode), len(goCode))
					}
				}
			}
		})
	}
}