// 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"
	"runtime"
	"strings"
	"testing"

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

// The differential kernels for the DATA-path and front-end gaps are kept in
// testdata/verify beside the campaign's other kernels; the verify package's
// suites are not open to the asm package, so this test is their runner: each
// kernel assembles through gasm and through go tool asm, and the functions'
// bytes must agree with the relocation sites masked on both sides.

// toolAsmObject assembles path with the installed toolchain's assembler for
// goarch ("" = the host) and returns the object bytes.
func toolAsmObject(t *testing.T, path, goarch string) []byte {
	t.Helper()
	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")

	pkg := strings.TrimSuffix(filepath.Base(path), ".s")
	pkg = strings.TrimSuffix(pkg, "_amd64")
	pkg = strings.TrimSuffix(pkg, "_arm64")

	objPath := filepath.Join(t.TempDir(), "oracle.o")
	cmd := exec.Command(goBin, "tool", "asm", "-I", includeDir, "-p", pkg, "-o", objPath, path)
	if goarch != "" {
		environ := os.Environ()
		env := make([]string, 0, len(environ)+1)
		for _, e := range environ {
			if !strings.HasPrefix(e, "GOARCH=") {
				env = append(env, e)
			}
		}
		cmd.Env = append(env, "GOARCH="+goarch)
	}
	if out, err := cmd.CombinedOutput(); err != nil {
		t.Fatalf("go tool asm %s: %v\n%s", filepath.Base(path), err, out)
	}
	obj, err := os.ReadFile(objPath)
	if err != nil {
		t.Fatal(err)
	}
	return obj
}

// oracleFuncCode extracts the non-package TEXT functions' code bytes from a
// toolchain object, keyed by the name the object records (pkg.name).
func oracleFuncCode(t *testing.T, obj []byte) map[string][]byte {
	t.Helper()
	v := openGoobj(t, obj)
	le := binary.LittleEndian
	const symSize = 21
	nps := v.syms(blkNonpkgdef)
	data := v.blk(blkData)
	didx := v.blk(blkDataIdx)
	preceding := 0
	for _, bi := range []int{blkSymdef, blkHashed64def, blkHasheddef} {
		preceding += len(v.blk(bi)) / symSize
	}
	total := preceding + len(nps)
	out := make(map[string][]byte, len(nps))
	for i, s := range nps {
		if s.typ != kindSTEXT {
			continue
		}
		start := le.Uint32(didx[4*(preceding+i):])
		end := uint32(len(data))
		if preceding+i+1 < total {
			end = le.Uint32(didx[4*(preceding+i+1):])
		}
		out[s.name] = data[start:end]
	}
	return out
}

// maskCode zeroes every relocation field, the way the toolchain's object
// leaves them for the linker.
func maskCode(code []byte, relocs []Reloc) []byte {
	for _, r := range relocs {
		for j := r.Off; j < r.Off+4 && j < len(code); j++ {
			code[j] = 0
		}
	}
	return code
}

// code assembles src for amd64 and returns the image's code bytes.
func code(path, src string) []byte {
	f, errs := parser.Parse(path, src)
	if len(errs) > 0 {
		return nil
	}
	img, err := AssembleFile(f)
	if err != nil {
		return nil
	}
	return img.Code
}

// TestDifferentialKernels pins the new kernels against the oracle.
func TestDifferentialKernels(t *testing.T) {
	if runtime.GOARCH != "amd64" {
		t.Skip("the amd64 kernels assume an amd64 host assembler default")
	}
	for _, k := range []struct {
		path   string
		goarch string
		arm64  bool
	}{
		{filepath.Join("..", "testdata", "verify", "datarel_amd64.s"), "", false},
		{filepath.Join("..", "testdata", "verify", "divslash_amd64.s"), "", false},
		{filepath.Join("..", "testdata", "verify", "semicolons_amd64.s"), "", false},
		{filepath.Join("..", "testdata", "verify", "datarel_arm64.s"), "arm64", true},
		{filepath.Join("..", "testdata", "verify", "divslash_arm64.s"), "arm64", true},
	} {
		t.Run(filepath.Base(k.path), func(t *testing.T) {
			src, err := os.ReadFile(k.path)
			if err != nil {
				t.Fatalf("read: %v", err)
			}
			f, errs := parser.Parse(k.path, string(src))
			if len(errs) > 0 {
				t.Fatalf("parse: %v", errs)
			}
			var img *Image
			if k.arm64 {
				img, err = AssembleFileARM64(f)
			} else {
				img, err = AssembleFile(f)
			}
			if err != nil {
				t.Fatalf("assemble: %v", err)
			}
			gt := oracleFuncCode(t, toolAsmObject(t, k.path, k.goarch))
			// The oracle keys its functions by the qualified object name
			// (pkg.name); match on the local part.
			byLocal := make(map[string][]byte, len(gt))
			for name, code := range gt {
				if _, after, ok := strings.Cut(name, "."); ok {
					name = after
				}
				byLocal[name] = code
			}

			matched := 0
			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 {
					t.Errorf("%s: not in ground truth (%d functions: %v)", fn.Name, len(gt), keysOf(byLocal))
					continue
				}
				goCode = maskCode(append([]byte(nil), goCode...), fn.Relocs)
				cmpLen := min(len(goCode), len(gasmCode))
				if !bytes.Equal(gasmCode[:cmpLen], goCode[:cmpLen]) {
					t.Errorf("%s: MISMATCH gasm=%d go=%d bytes\ngasm %x\ngo   %x", fn.Name, len(gasmCode), len(goCode), gasmCode, goCode)
					continue
				}
				for _, b := range goCode[len(gasmCode):] {
					if b != 0 {
						t.Errorf("%s: non-zero trailing bytes in go tool asm output", fn.Name)
						break
					}
				}
				matched++
				t.Logf("%s: MATCH (%d bytes)", fn.Name, len(gasmCode))
			}
			if matched == 0 {
				t.Fatal("no functions matched")
			}
		})
	}
}

func keysOf(m map[string][]byte) []string {
	out := make([]string, 0, len(m))
	for k := range m {
		out = append(out, k)
	}
	return out
}

// TestSemicolonSpellingParity pins that the ';' statement separator changes
// nothing about the encoding: the one-line spelling assembles to exactly the
// bytes of the same statements written one per line.
func TestSemicolonSpellingParity(t *testing.T) {
	for _, tt := range []struct{ one, two string }{
		{"\tROLQ $3, DI; ROLQ $13, DI\n", "\tROLQ $3, DI\n\tROLQ $13, DI\n"},
		{"\tREP; MOVSQ\n", "\tREP\n\tMOVSQ\n"},
		{"\tXORQ AX, AX; XORQ CX, CX\n", "\tXORQ AX, AX\n\tXORQ CX, CX\n"},
	} {
		one := code("t.s", "TEXT \u00b7f(SB), NOSPLIT, $0\n"+tt.one+"\tRET\n")
		two := code("t.s", "TEXT \u00b7f(SB), NOSPLIT, $0\n"+tt.two+"\tRET\n")
		if !bytes.Equal(one, two) {
			t.Errorf("semicolon spelling %q: %x, want the two-line bytes %x", tt.one, one, two)
		}
	}
}