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

package asm

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

	"sourcedock.dev/petrbalvin/gasm-devkit/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")
	}
}

// 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
`
	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 main() {
	if add(20, 22) != 42 {
		panic("bad add")
	}
}
`
	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 pkgArch, work, linkLine, asmObj string
	for _, line := range strings.Split(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, "pack r") && strings.Contains(line, "_pkg_.a"):
			pkgArch = strings.TrimSpace(strings.SplitN(line, "pack r", 2)[1])
			pkgArch = strings.Fields(strings.SplitN(pkgArch, "#", 2)[0])[0]
		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")
	}
}