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

package asm

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

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

// TestGOObjectAARCH64Structure checks the basic structure of the emitted
// AArch64 GOOBJ: the preamble, the magic, the block offsets and the
// non-package symbol definitions.
func TestGOObjectAARCH64Structure(t *testing.T) {
	f, errs := parser.Parse("k_arm64.s", `
#include "textflag.h"

TEXT ·add(SB), NOSPLIT, $0-24
	MOVD a+0(FP), R4
	MOVD b+8(FP), R5
	ADD R5, R4, R4
	MOVD R4, ret+16(FP)
	RET
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("AssembleFileARM64: %v", err)
	}
	obj, err := img.GOObjectAARCH64("testpkg", "k_arm64.s")
	if err != nil {
		t.Fatalf("GOObjectAARCH64: %v", err)
	}

	// Check preamble.
	idx := strings.Index(string(obj), "\n!\n")
	if idx < 0 {
		t.Fatal("missing preamble separator")
	}
	preamble := string(obj[:idx])
	if !strings.HasPrefix(preamble, "go object") {
		t.Errorf("preamble = %q, want 'go object ...'", preamble)
	}

	// Check GOOBJ magic.
	magicIdx := idx + 3
	if magicIdx+8 > len(obj) || string(obj[magicIdx:magicIdx+8]) != "\x00go120ld" {
		t.Error("missing GOOBJ magic")
	}

	// The object should contain the function's code.
	if len(img.Code) == 0 {
		t.Error("no code generated")
	}
}

// TestGOObjectAARCH64PairReloc pins the ADRP-pair relocation shape against
// the toolchain's own object for the same source: exactly one R_ADDRARM64
// of Siz 8 at the ADRP word (cmd/internal/obj/arm64/asm7.go adds a single
// Siz-8 relocation per pair and the linker patches both instructions from
// it).  gasm's assembler records the ADRP+ADD form as two word relocs; the
// emitter must coalesce them, not emit two Siz-4 records.
func TestGOObjectAARCH64PairReloc(t *testing.T) {
	f, errs := parser.Parse("gv_arm64.s", `
#include "textflag.h"

TEXT ·getv(SB), NOSPLIT, $0-8
	MOVD $v<>(SB), R4
	MOVD R4, ret+0(FP)
	RET

GLOBL v<>(SB), RODATA, $8
DATA v<>+0(SB)/8, $7
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("AssembleFileARM64: %v", err)
	}
	obj, err := img.GOObjectAARCH64("main", "gv_arm64.s")
	if err != nil {
		t.Fatalf("GOObjectAARCH64: %v", err)
	}
	v := openGoobj(t, obj)
	relocs := v.blk(blkReloc)
	le := binary.LittleEndian
	// Two DWARF relocs on the lines/DIE symbols, then the code's one pair
	// relocation.
	if len(relocs) != 3*23 {
		t.Fatalf("relocs = %d bytes, want three entries", len(relocs))
	}
	cr := relocs[2*23:]
	if off := int32(le.Uint32(cr[0:])); off != 0 {
		t.Errorf("pair reloc off = %d, want 0 (the ADRP word)", off)
	}
	if siz := cr[4]; siz != 8 {
		t.Errorf("pair reloc siz = %d, want 8", siz)
	}
	if typ := le.Uint16(cr[5:]); typ != relocArm64Addr {
		t.Errorf("pair reloc type = %d, want %d (R_ADDRARM64)", typ, relocArm64Addr)
	}
	if pkg := le.Uint32(cr[15:]); pkg != pkgIdxSelf {
		t.Errorf("pair reloc PkgIdx = %#x, want pkgIdxSelf", pkg)
	}
	// The GLOBL is the first package definition.
	if sym := le.Uint32(cr[19:]); sym != 0 {
		t.Errorf("pair reloc SymIdx = %d, want 0 (the GLOBL definition)", sym)
	}
}

// TestGOObjectAARCH64Link does an end-to-end link test: it cross-compiles a
// Go program for arm64, substitutes the gasm-produced object into the package
// archive, re-links with cmd/link, and verifies the symbol appears in the
// resulting binary.  The binary is not executed (no arm64 host or qemu).
// Skipped when no Go toolchain is available.
func TestGOObjectAARCH64Link(t *testing.T) {
	goBin, err := exec.LookPath("go")
	if err != nil {
		t.Skip("no Go toolchain available")
	}
	dir := t.TempDir()
	asmSrc := `#include "textflag.h"
TEXT ·add(SB), NOSPLIT, $0-24
	MOVD a+0(FP), R4
	MOVD b+8(FP), R5
	ADD R5, R4, R4
	MOVD R4, ret+16(FP)
	RET

TEXT ·getv(SB), NOSPLIT, $0-8
	MOVD $v<>(SB), R4
	MOVD R4, ret+0(FP)
	RET

GLOBL v<>(SB), RODATA, $8
DATA v<>+0(SB)/8, $7
`
	if err := os.WriteFile(filepath.Join(dir, "main_arm64.s"), []byte(asmSrc), 0o644); err != nil {
		t.Fatal(err)
	}
	mainSrc := `package main

func add(a, b int64) int64
func getv() *int64

func main() {
	if add(20, 22) != 42 {
		panic("bad add")
	}
	if getv() == nil {
		panic("bad getv")
	}
}
`
	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 a64link\n\ngo 1.21\n"), 0o644); err != nil {
		t.Fatal(err)
	}

	// Capture the cross build (GOARCH=arm64): the package archive and the
	// link line.
	build := exec.Command(goBin, "build", "-x", "-work", "-o", filepath.Join(dir, "prog"), ".")
	build.Dir = dir
	build.Env = append(os.Environ(), "GOARCH=arm64")
	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_arm64.s") && !strings.Contains(line, "-gensymabis"):
			asmObj = fieldAfter(line, "-o")
		case strings.Contains(line, "/link ") && strings.Contains(line, "-importcfg"):
			linkLine = line
		}
	}
	if work == "" || asmObj == "" {
		t.Skipf("could not parse build log (work=%q asmObj=%q)", work, asmObj)
	}
	defer os.RemoveAll(work)

	// Expand $WORK in the object path.
	asmObj = strings.ReplaceAll(asmObj, "$WORK", work)

	// Read the toolchain-produced object and assemble the same source with gasm.
	src, err := os.ReadFile(filepath.Join(dir, "main_arm64.s"))
	if err != nil {
		t.Fatal(err)
	}
	f, errs := parser.Parse("main_arm64.s", string(src))
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("AssembleFileARM64: %v", err)
	}
	gasmObj, err := img.GOObjectAARCH64("a64link", "main_arm64.s")
	if err != nil {
		t.Fatalf("GOObjectAARCH64: %v", err)
	}

	// Replace the toolchain-produced object with gasm's.
	if err := os.WriteFile(asmObj, gasmObj, 0o644); err != nil {
		t.Fatalf("write gasm object: %v", err)
	}

	// Re-link.
	if linkLine == "" {
		t.Skip("could not find link command in build log")
	}
	// Expand $WORK in the link command.
	linkLine = strings.ReplaceAll(linkLine, "$WORK", work)
	linkCmd := exec.Command("bash", "-c", "cd "+dir+" && "+linkLine)
	linkCmd.Env = append(os.Environ(), "GOARCH=arm64")
	if out, err := linkCmd.CombinedOutput(); err != nil {
		t.Fatalf("re-link with gasm object: %v\n%s", err, out)
	}

	// Verify the binary exists and contains the symbol.
	binPath := filepath.Join(dir, "prog")
	if _, err := os.Stat(binPath); err != nil {
		t.Fatalf("binary not found: %v", err)
	}
	binData, err := os.ReadFile(binPath)
	if err != nil {
		t.Fatalf("read binary: %v", err)
	}
	if !strings.Contains(string(binData), "add") && !strings.Contains(string(binData), "a64link") {
		t.Error("binary does not contain expected symbol")
	}
}

// TestGOObjectAARCH64DataSymbolLink does for symbol-valued DATA fields what
// the rt0 files do ("DATA _rt0…lib+0(SB)/8, $_rt0…lib(SB)"): the gasm object
// carries an R_ADDR against the file's own TEXT symbol, the toolchain links
// it, and the binary is checked for the symbol (no arm64 host to run it).
func TestGOObjectAARCH64DataSymbolLink(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
	MOVD entry+0(SB), R4
	MOVD R4, ret+0(FP)
	RET
`
	if err := os.WriteFile(filepath.Join(dir, "main_arm64.s"), []byte(asmSrc), 0o644); err != nil {
		t.Fatal(err)
	}
	mainSrc := `package main

func keepme()
func entryptr() uintptr

func main() {
	if entryptr() == 0 {
		panic("the entry word is empty")
	}
}
`
	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 a64dlink\n\ngo 1.21\n"), 0o644); err != nil {
		t.Fatal(err)
	}

	build := exec.Command(goBin, "build", "-x", "-work", "-o", filepath.Join(dir, "prog"), ".")
	build.Dir = dir
	build.Env = append(os.Environ(), "GOARCH=arm64")
	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_arm64.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)

	src, err := os.ReadFile(filepath.Join(dir, "main_arm64.s"))
	if err != nil {
		t.Fatal(err)
	}
	f, errs := parser.Parse("main_arm64.s", string(src))
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("AssembleFileARM64: %v", err)
	}
	gasmObj, err := img.GOObjectAARCH64("a64dlink", "main_arm64.s")
	if err != nil {
		t.Fatalf("GOObjectAARCH64: %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)
	linkCmd.Env = append(os.Environ(), "GOARCH=arm64")
	if out, err := linkCmd.CombinedOutput(); err != nil {
		t.Fatalf("re-link with gasm object: %v\n%s", err, out)
	}
	binData, err := os.ReadFile(filepath.Join(dir, "prog"))
	if err != nil {
		t.Fatal(err)
	}
	if !strings.Contains(string(binData), "keepme") {
		t.Error("binary does not contain the keepme symbol")
	}
}