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

package asm

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

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

// TestAssembleFileStaticData checks the whole-image layout; code, padding
// and the data section; and that the RIP-relative displacements of static
// symbol loads resolve to the right bytes.
func TestAssembleFileStaticData(t *testing.T) {
	f, errs := parser.Parse("d_amd64.s", `
#include "textflag.h"
TEXT ·load(SB), NOSPLIT, $0
	VMOVDQU mask<>(SB), X15
	MOVL small<>(SB), AX
	RET
GLOBL mask<>(SB), RODATA, $16
DATA mask<>+0(SB)/4, $0x80020100
DATA mask<>+4(SB)/4, $0x80050403
DATA mask<>+8(SB)/4, $0x80080706
DATA mask<>+12(SB)/4, $0x800B0A09
GLOBL small<>(SB), RODATA, $4
DATA small<>+0(SB)/4, $0x1234
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("AssembleFile: %v", err)
	}

	// Code (15 bytes) + 1 pad byte to align the data section to 16:
	//   VMOVDQU mask<>(SB), X15   c5 7a 6f 3d 08 00 00 00   (disp = 16 − 8)
	//   MOVL small<>(SB), AX      8b 05 12 00 00 00         (disp = 32 − 14)
	//   RET                       c3
	// Data: pad, mask (16 bytes), small (4 bytes).
	want := "c57a6f3d080000008b0512000000c300" +
		"000102800304058006070880090a0b80" +
		"34120000"
	if got := strings.ReplaceAll(hexBytes(img.Bytes()), " ", ""); got != want {
		t.Errorf("image bytes:\n got  %s\n want %s", got, want)
	}
	if img.Symbols["mask"] != 16 || img.Symbols["small"] != 32 {
		t.Errorf("symbol offsets = %v, want mask=16 small=32", img.Symbols)
	}
	if len(img.Funcs) != 1 || img.Funcs[0].Name != "load" || img.Funcs[0].Size != 15 {
		t.Errorf("funcs = %+v", img.Funcs)
	}
}

// TestAssembleFileErrors checks the static-symbol error paths.
func TestAssembleFileErrors(t *testing.T) {
	cases := []struct {
		name string
		src  string
		want string // substring of the error
	}{
		{
			"undefined symbol",
			`
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	VMOVDQU nope<>(SB), X0
	RET
`,
			"undefined symbol",
		},
		{
			"DATA without GLOBL",
			`
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	RET
DATA orphan<>+0(SB)/4, $1
`,
			"no matching GLOBL",
		},
		{
			"DATA exceeds size",
			`
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	RET
GLOBL tiny<>(SB), RODATA, $4
DATA tiny<>+0(SB)/8, $1
`,
			"exceeds GLOBL size",
		},
		{
			"DATA bad width",
			`
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	RET
GLOBL odd<>(SB), RODATA, $4
DATA odd<>+0(SB)/3, $1
`,
			"invalid width",
		},
	}
	for _, c := range cases {
		f, errs := parser.Parse("e_amd64.s", c.src)
		if len(errs) > 0 {
			t.Fatalf("%s: parse: %v", c.name, errs)
		}
		if _, err := AssembleFile(f); err == nil || !strings.Contains(err.Error(), c.want) {
			t.Errorf("%s: error %v, want substring %q", c.name, err, c.want)
		}
	}

	// A static-symbol operand is unresolvable in single-function assembly.
	fn := firstText(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVQ x<>(SB), AX
	RET
GLOBL x<>(SB), RODATA, $8
DATA x<>+0(SB)/4, $1
`)
	if _, _, err := Assemble(fn); err == nil || !strings.Contains(err.Error(), "file-level assembly") {
		t.Errorf("single-function SB: error %v, want a file-level-assembly error", err)
	}
}

// TestCollectDataNumericFlags pins the numeric GLOBL flag constants from
// runtime/textflag.h: DUPOK is 2, RODATA is 8, and combinations arrive as
// one number (9 = NOPROF|RODATA, 10 = RODATA|DUPOK).
func TestCollectDataNumericFlags(t *testing.T) {
	tests := []struct {
		flags  string
		rodata bool
		dupok  bool
	}{
		{"2", false, true},
		{"8", true, false},
		{"9", true, false}, // NOPROF|RODATA, not DUPOK
		{"10", true, true}, // RODATA|DUPOK
		{"RODATA", true, false},
		{"DUPOK", false, true},
		{"RODATA|DUPOK", true, true},
	}
	for _, tt := range tests {
		src := "TEXT \u00b7f(SB), NOSPLIT, $0\n\tRET\nGLOBL sym(SB), " + tt.flags + ", $8\n"
		f, errs := parser.Parse("f_amd64.s", src)
		if len(errs) > 0 {
			t.Fatalf("parse %q: %v", tt.flags, errs)
		}
		img, err := AssembleFile(f)
		if err != nil {
			t.Fatalf("assemble %q: %v", tt.flags, err)
		}
		if len(img.DataSyms) != 1 {
			t.Fatalf("%q: data syms = %d, want 1", tt.flags, len(img.DataSyms))
		}
		d := img.DataSyms[0]
		if d.Rodata != tt.rodata || d.Dupok != tt.dupok {
			t.Errorf("flags %q: rodata=%v dupok=%v, want rodata=%v dupok=%v",
				tt.flags, d.Rodata, d.Dupok, tt.rodata, tt.dupok)
		}
	}
}

// TestCollectDataSymbolValue covers the symbol-valued DATA field ("DATA
// s+0(SB)/8, $other(SB)", the rt0 spelling): the field stays zero in the
// image and the relocation is recorded against the named symbol, whatever
// the file defines (a TEXT function, a GLOBL) or leaves external.
func TestCollectDataSymbolValue(t *testing.T) {
	src := `#include "textflag.h"
TEXT ·Keep(SB), NOSPLIT, $0-8
	MOVQ target+0(FP), AX
	RET
GLOBL holder(SB), NOPTR, $32
DATA holder+0(SB)/8, $·Keep(SB)
DATA holder+8(SB)/8, $·Keep+5(SB)
DATA holder+16(SB)/8, $holder(SB)
GLOBL spare(SB), NOPTR, $8
DATA spare+0(SB)/8, $extvar(SB)
`
	f, errs := parser.Parse("f_amd64.s", src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("assemble: %v", err)
	}
	byName := map[string]DataSymbol{}
	for _, d := range img.DataSyms {
		byName[d.Name] = d
	}
	want := []struct {
		sym    string
		off    int
		name   string
		addend int64
		ext    bool
	}{
		{"holder", 0, "Keep", 0, false},
		{"holder", 8, "Keep", 5, false},
		{"holder", 16, "holder", 0, false},
		{"spare", 0, "extvar", 0, true},
	}
	var flat []struct {
		sym string
		r   Reloc
	}
	for _, d := range img.DataSyms {
		for _, r := range d.Relocs {
			flat = append(flat, struct {
				sym string
				r   Reloc
			}{d.Name, r})
		}
	}
	if len(flat) != len(want) {
		t.Fatalf("data relocations = %d, want %d", len(flat), len(want))
	}
	for i, w := range want {
		g := flat[i]
		r := g.r
		if g.sym != w.sym {
			t.Errorf("relocation %d sits on %q, want %q", i, g.sym, w.sym)
			continue
		}
		if r.Off != w.off || r.Name != w.name || r.Addend != w.addend || r.External != w.ext {
			t.Errorf("relocation %d = {+%d %q addend %d ext %v}, want {+%d %q addend %d ext %v}",
				i, r.Off, r.Name, r.Addend, r.External, w.off, w.name, w.addend, w.ext)
		}
		if r.Kind != RelAddr {
			t.Errorf("relocation %d kind = %v, want RelAddr", i, r.Kind)
		}
		if r.Siz != 8 {
			t.Errorf("relocation %d siz = %d, want 8", i, r.Siz)
		}
	}
	// The fields themselves stay zero: only the linker fills them.
	for _, b := range img.Data {
		if b != 0 {
			t.Fatal("data section is not all zero before relocation")
		}
	}
	if len(img.Externals) != 1 || img.Externals[0] != "extvar" {
		t.Errorf("Externals = %v, want [extvar]", img.Externals)
	}
}

// TestGOObjectDataSymbolReloc pins the GOOBJ record a symbol-valued DATA
// field produces, against the shape the toolchain emits for the same
// source: an R_ADDR of the DATA width at the field offset, pkgIdxNone plus
// the non-package definition index when the target is the file's own TEXT
// function (the rt0 lib entry spelling).
func TestGOObjectDataSymbolReloc(t *testing.T) {
	f, errs := parser.Parse("f_amd64.s", `#include "textflag.h"
TEXT ·Keep(SB), NOSPLIT, $0-8
	RET
GLOBL holder(SB), NOPTR, $16
DATA holder+0(SB)/8, $·Keep+5(SB)
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("assemble: %v", err)
	}
	obj, err := img.GOObject("main", "f_amd64.s")
	if err != nil {
		t.Fatalf("GOObject: %v", err)
	}
	v := openGoobj(t, obj)
	// Walk every relocation record; the data record is the one of Siz 8
	// and type R_ADDR.
	var off, add int64
	var pkg, sym uint32
	found := false
	for data := v.blk(blkReloc); len(data) >= 23; data = data[23:] {
		if data[4] != 8 || binary.LittleEndian.Uint16(data[5:]) != relocAddr {
			continue
		}
		found = true
		off = int64(int32(binary.LittleEndian.Uint32(data[0:])))
		add = int64(binary.LittleEndian.Uint64(data[7:]))
		pkg = binary.LittleEndian.Uint32(data[15:])
		sym = binary.LittleEndian.Uint32(data[19:])
		break
	}
	if !found {
		t.Fatal("no data relocation record in the object")
	}
	if off != 0 || add != 5 {
		t.Errorf("data reloc = {off %d addend %d}, want {off 0 addend 5}", off, add)
	}
	if pkg != pkgIdxNone {
		t.Errorf("data reloc pkg = %#x, want pkgIdxNone (the TEXT function)", pkg)
	}
	// The function's non-package definition index: the four pc tables
	// precede it, so index 4.
	if sym != 4 {
		t.Errorf("data reloc sym = %d, want 4", sym)
	}
}

// TestGOObjectDataSymbolLink is the end-to-end proof for symbol-valued DATA
// fields: the gasm object is substituted for the toolchain's and re-linked,
// then executed, and the linked data word must hold the real address of the
// function the DATA line named (runtime.FuncForPC identifies it).
func TestGOObjectDataSymbolLink(t *testing.T) {
	goBin, err := exec.LookPath("go")
	if err != nil {
		t.Skip("no Go toolchain available")
	}
	dir := t.TempDir()
	asmSrc := `#include "textflag.h"
GLOBL entry(SB), NOPTR, $8
DATA entry+0(SB)/8, $·keepme(SB)

TEXT ·keepme(SB), NOSPLIT, $0-0
	RET

TEXT ·entryptr(SB), NOSPLIT, $0-8
	MOVQ entry+0(SB), AX
	MOVQ AX, ret+0(FP)
	RET
`
	if err := os.WriteFile(filepath.Join(dir, "main_amd64.s"), []byte(asmSrc), 0o644); err != nil {
		t.Fatal(err)
	}
	mainSrc := `package main

import "runtime"

func keepme()
func entryptr() uintptr

func main() {
	pc := entryptr()
	fn := runtime.FuncForPC(pc)
	if fn == nil {
		panic("the entry word does not point at a function")
	}
	if fn.Name() != "main.keepme" {
		panic("the entry word points at " + fn.Name())
	}
}
`
	if err := os.WriteFile(filepath.Join(dir, "main.go"), []byte(mainSrc), 0o644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(filepath.Join(dir, "go.mod"), []byte("module dlink\n\ngo 1.21\n"), 0o644); err != nil {
		t.Fatal(err)
	}

	// Capture the build: the package archive's asm object and the link line.
	build := exec.Command(goBin, "build", "-x", "-work", "-o", filepath.Join(dir, "prog"), ".")
	build.Dir = dir
	buildLog, err := build.CombinedOutput()
	if err != nil {
		t.Fatalf("baseline build: %v\n%s", err, buildLog)
	}
	var work, linkLine, asmObj string
	for line := range strings.SplitSeq(string(buildLog), "\n") {
		switch {
		case strings.HasPrefix(line, "WORK="):
			work = strings.TrimPrefix(line, "WORK=")
		case strings.Contains(line, "/asm ") && strings.Contains(line, "main_amd64.s") && !strings.Contains(line, "-gensymabis"):
			asmObj = fieldAfter(line, "-o")
		case strings.Contains(line, "/link ") && strings.Contains(line, "-importcfg"):
			linkLine = line
		}
	}
	if work == "" || asmObj == "" || linkLine == "" {
		t.Skipf("could not parse build log (work=%q asmObj=%q link=%q)", work, asmObj, linkLine)
	}
	defer os.RemoveAll(work)
	asmObj = strings.ReplaceAll(asmObj, "$WORK", work)
	linkLine = strings.ReplaceAll(linkLine, "$WORK", work)

	// Assemble the same source with gasm and substitute the object.
	src, err := os.ReadFile(filepath.Join(dir, "main_amd64.s"))
	if err != nil {
		t.Fatal(err)
	}
	f, errs := parser.Parse("main_amd64.s", string(src))
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("AssembleFile: %v", err)
	}
	gasmObj, err := img.GOObject("dlink", "main_amd64.s")
	if err != nil {
		t.Fatalf("GOObject: %v", err)
	}
	if err := os.WriteFile(asmObj, gasmObj, 0o644); err != nil {
		t.Fatalf("write gasm object: %v", err)
	}
	linkCmd := exec.Command("bash", "-c", "cd "+dir+" && "+linkLine)
	if out, err := linkCmd.CombinedOutput(); err != nil {
		t.Fatalf("re-link with gasm object: %v\n%s", err, out)
	}

	// The linked program must run and find the right function behind the
	// data word.
	out, err := exec.Command(filepath.Join(dir, "prog")).CombinedOutput()
	if err != nil {
		t.Fatalf("linked program failed: %v\n%s", err, out)
	}
}

// TestCollectDataFloatAndStringValues covers the non-integer DATA values the
// runtime's math and asm files use: floating-point initialisers store their
// IEEE-754 bits (/4 the float32 rounding, /8 the full double) and string
// initialisers write their bytes zero-padded within the declared width.
func TestCollectDataFloatAndStringValues(t *testing.T) {
	src := `#include "textflag.h"
TEXT ·Keep(SB), NOSPLIT, $0-8
	RET
GLOBL vals<>(SB), RODATA, $44
DATA vals<>+0(SB)/8, $0.5
DATA vals<>+8(SB)/8, $-1.0
DATA vals<>+16(SB)/4, $1.5
DATA vals<>+20(SB)/16, $"call frame too "
DATA vals<>+36(SB)/4, $"hi"
`
	f, errs := parser.Parse("fvals_amd64.s", src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("AssembleFile: %v", err)
	}
	byName := map[string]DataSymbol{}
	for _, d := range img.DataSyms {
		byName[d.Name] = d
	}
	d := byName["vals"]
	if d.Size != 44 {
		t.Fatalf("vals size = %d, want 44", d.Size)
	}
	buf := img.Data[d.Offset : d.Offset+44]
	// 0.5 = 0x3FE0000000000000, -1.0 = 0xBFF0000000000000 (float64);
	// 1.5 = 0x3FC00000 (float32).
	for _, c := range []struct {
		off  int
		want []byte
	}{
		{0, []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xE0, 0x3F}},
		{8, []byte{0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xF0, 0xBF}},
		{16, []byte{0x00, 0x00, 0xC0, 0x3F}},
		{20, []byte("call frame too ")},
		{36, []byte{'h', 'i', 0x00, 0x00}},
	} {
		if string(buf[c.off:c.off+len(c.want)]) != string(c.want) {
			t.Errorf("vals+%d: got % x, want % x", c.off, buf[c.off:c.off+len(c.want)], c.want)
		}
	}
}

// TestCollectDataValueErrors pins the value-kind width rules: a float needs
// width 4 or 8, a string must fit its declared width, and a bad float
// literal is diagnosed rather than stored.
func TestCollectDataValueErrors(t *testing.T) {
	cases := []string{
		`GLOBL v<>(SB), RODATA, $4
DATA v<>+0(SB)/1, $0.5`,
		`GLOBL v<>(SB), RODATA, $2
DATA v<>+0(SB)/2, $"toolarge"`,
	}
	for i, src := range cases {
		full := "#include \"textflag.h\"\nTEXT ·Keep(SB), NOSPLIT, $0-8\n\tRET\n" + src
		f, errs := parser.Parse(fmt.Sprintf("verr%d_amd64.s", i), full)
		if len(errs) > 0 {
			t.Fatalf("case %d parse: %v", i, errs)
		}
		if _, err := AssembleFile(f); err == nil {
			t.Errorf("case %d: expected an error, got none", i)
		}
	}
}