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

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

// TestGOObjectLOONG64Structure checks the emitted loong64 object's blocks:
// the symbol tables, the function code bytes and the relocation wiring.
func TestGOObjectLOONG64Structure(t *testing.T) {
	f, errs := parser.Parse("k_loong64.s", `
#include "textflag.h"

TEXT ·add(SB), NOSPLIT, $0-24
	MOVV a+0(FP), R4
	MOVV b+8(FP), R5
	ADDV R5, R4, R4
	MOVV R4, ret+16(FP)
	RET

GLOBL ·table<>(SB), RODATA, $8
DATA ·table<>+0(SB)/8, $0x1122334455667788
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileLOONG64(f)
	if err != nil {
		t.Fatalf("AssembleFileLOONG64: %v", err)
	}
	obj, err := img.GOObjectLOONG64("testpkg", "k_loong64.s")
	if err != nil {
		t.Fatalf("GOObjectLOONG64: %v", err)
	}
	v := openGoobj(t, obj)

	// Package defs: the static GLOBL, then the FuncInfo and the two DWARF
	// symbols (debug_line program, subprogram DIE).
	defs := v.syms(blkSymdef)
	if len(defs) != 4 {
		t.Fatalf("symdefs = %d, want 4", len(defs))
	}
	if defs[0].name != "table" || defs[0].abi != 0xffff || defs[0].typ != kindSRODATA || defs[0].size != 8 {
		t.Errorf("table symbol = %+v", defs[0])
	}
	if defs[1].name != "" || defs[1].typ != kindSDATA || defs[1].size != 28 {
		t.Errorf("funcinfo symbol = %+v", defs[1])
	}
	if defs[2].name != "" || defs[2].typ != kindSDWARFLINES || defs[2].size == 0 {
		t.Errorf("lines symbol = %+v", defs[2])
	}
	if defs[3].name != "" || defs[3].typ != kindSDWARFFCN || defs[3].size == 0 {
		t.Errorf("DIE symbol = %+v", defs[3])
	}

	// Non-package defs: four pc tables and the function.
	nps := v.syms(blkNonpkgdef)
	if len(nps) != 5 {
		t.Fatalf("nonpkgdefs = %d, want 5", len(nps))
	}
	fn := nps[4]
	if fn.name != "testpkg.add" || fn.typ != kindSTEXT || fn.flag != symFlagNoSplit || fn.size != 20 {
		t.Errorf("add symbol = %+v", fn)
	}

	// The function code: 20 bytes, the ground-truth encoding.  It sits
	// after the GLOBL, FuncInfo, two DWARF symbols and four pc tables.
	dataIdx := v.blk(blkDataIdx)
	dataBlk := v.blk(blkData)
	le := binary.LittleEndian
	dOff := le.Uint32(dataIdx[8*4:])
	code := dataBlk[dOff : dOff+20]
	want := []byte{
		0x64, 0x20, 0xc0, 0x28, // ld.d r4, 8(r3)
		0x65, 0x40, 0xc0, 0x28, // ld.d r5, 16(r3)
		0x84, 0x94, 0x10, 0x00, // add.d r4, r4, r5
		0x64, 0x60, 0xc0, 0x29, // st.d r4, 24(r3)
		0x20, 0x00, 0x00, 0x4c, // jirl r0, r1, 0
	}
	for i := range want {
		if code[i] != want[i] {
			t.Fatalf("code byte %d = %02x, want %02x", i, code[i], want[i])
		}
	}

	// The debug_line program: LNE_set_address (the R_ADDR relocation
	// carries the function address), then one row per line change; the
	// TEXT is on line 4 (a leading blank line precedes the include), the
	// instructions on lines 5-9; an advance to the 20-byte end and an
	// end-of-sequence.
	linesOff := le.Uint32(dataIdx[4*2:])
	lines := dataBlk[linesOff : linesOff+21]
	wantLines := []byte{
		0x00, 0x09, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, // LNE_set_address
		0x13,       // pc 0, line 5
		0x38,       // pc 4, line 6
		0x38,       // pc 8, line 7
		0x38,       // pc 12, line 8
		0x38,       // pc 16, line 9
		0x02, 0x04, // advance_pc to 20
		0x00, 0x01, 0x01, // end_sequence
	}
	for i := range wantLines {
		if lines[i] != wantLines[i] {
			t.Fatalf("lines byte %d = %02x, want %02x", i, lines[i], wantLines[i])
		}
	}

	// The subprogram DIE: abbrev 3 (FUNCTION), the qualified name, the
	// addrx low_pc slot (R_DWTXTADDR_U4), the size as high_pc, the
	// call-frame-CFA frame base, decl file/line and the external flag.
	dieOff := le.Uint32(dataIdx[4*3:])
	die := dataBlk[dieOff : dieOff+27]
	wantDie := []byte{
		0x03,
		't', 'e', 's', 't', 'p', 'k', 'g', '.', 'a', 'd', 'd', 0,
		0x00, 0x00, 0x00, 0x00, // low_pc: addrx slot
		0x14,       // high_pc: 20
		0x01, 0x9c, // frame_base: DW_OP_call_frame_cfa
		0x01, 0x00, 0x00, 0x00, // decl_file: 1
		0x04, // decl_line: 4
		0x01, // external
		0x00, // end of children
	}
	for i := range wantDie {
		if die[i] != wantDie[i] {
			t.Fatalf("DIE byte %d = %02x, want %02x", i, die[i], wantDie[i])
		}
	}

	// The DWARF symbols carry the function-address references: R_ADDR for
	// the line program's set_address, R_DWTXTADDR_U4 for the DIE's addrx
	// slot, both against the function's non-package index.  The reloc
	// index counts relocations, not bytes.
	relocIdx := v.blk(blkRelocIdx)
	relocs := v.blk(blkReloc)
	if le.Uint32(relocIdx[4*2:]) != 0 || le.Uint32(relocIdx[4*3:]) != 1 || le.Uint32(relocIdx[4*4:]) != 2 {
		t.Fatalf("dwarf reloc index ranges: %d %d %d", le.Uint32(relocIdx[4*2:]), le.Uint32(relocIdx[4*3:]), le.Uint32(relocIdx[4*4:]))
	}
	lr := relocs[:23]
	if int32(le.Uint32(lr[0:])) != 3 || lr[4] != 8 || le.Uint16(lr[5:]) != relocAddr ||
		le.Uint32(lr[15:]) != pkgIdxNone || le.Uint32(lr[19:]) != 4 {
		t.Errorf("lines reloc = %x", lr)
	}
	dr := relocs[23:46]
	if int32(le.Uint32(dr[0:])) != 13 || dr[4] != 4 || le.Uint16(dr[5:]) != relocDWTXTADDRU4() ||
		le.Uint32(dr[15:]) != pkgIdxNone || le.Uint32(dr[19:]) != 4 {
		t.Errorf("die reloc = %x", dr)
	}

	// The pc-value deltas are in MinLC (4) units: the flat pcsp covers
	// the whole 20-byte function with a delta of 5.
	pcspOff := le.Uint32(dataIdx[4*4:])
	if got := dataBlk[pcspOff : pcspOff+3]; !bytes.Equal(got, []byte{0x02, 0x05, 0x00}) {
		t.Errorf("pcsp = %x, want 020500", got)
	}
}

// TestGOObjectLOONG64Link cross-compiles a Go program with the gasm-produced
// object substituted into the package archive, proving cmd/link accepts the
// emitted GOOBJ.  The binary is not executed (no LoongArch host or qemu).
// Skipped when no Go toolchain is available.
func TestGOObjectLOONG64Link(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
	MOVV a+0(FP), R4
	MOVV b+8(FP), R5
	ADDV R5, R4, R4
	MOVV R4, ret+16(FP)
	RET
`
	if err := os.WriteFile(filepath.Join(dir, "main_loong64.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 l64link\n\ngo 1.21\n"), 0o644); err != nil {
		t.Fatal(err)
	}

	// Capture the cross build (GOARCH=loong64): 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=loong64")
	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.SplitSeq(string(buildLog), "\n") {
		switch {
		case strings.HasPrefix(line, "WORK="):
			work = strings.TrimPrefix(line, "WORK=")
		case strings.Contains(line, "/asm ") && strings.Contains(line, "main_loong64.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 pkgArch == "" || linkLine == "" || asmObj == "" {
		t.Skip("could not locate the archive, asm output or link line in the build log")
	}
	pkgArch = strings.ReplaceAll(pkgArch, "$WORK", work)
	// The archive member holding the assembler's output is named after the
	// asm object file (main_loong64.o), as cmd/go packs it with `pack r`.
	asmMember := filepath.Base(strings.ReplaceAll(asmObj, "$WORK", work))

	// Assemble the same source with gasm and swap the object in.
	pf, perrs := parser.Parse(filepath.Join(dir, "main_loong64.s"), asmSrc)
	if len(perrs) > 0 {
		t.Fatalf("parse: %v", perrs)
	}
	pimg, err := AssembleFileLOONG64(pf)
	if err != nil {
		t.Fatalf("AssembleFileLOONG64: %v", err)
	}
	obj, err := pimg.GOObjectLOONG64("main", filepath.Join(dir, "main_loong64.s"))
	if err != nil {
		t.Fatalf("GOObjectLOONG64: %v", err)
	}

	// Extract the archive, substitute the object member, repack.
	membersDir := filepath.Join(dir, "members")
	if err := os.MkdirAll(membersDir, 0o755); err != nil {
		t.Fatal(err)
	}
	extract := exec.Command(goBin, "tool", "pack", "x", pkgArch)
	extract.Dir = membersDir
	extract.Env = append(os.Environ(), "GOARCH=loong64")
	if out, err := extract.CombinedOutput(); err != nil {
		t.Fatalf("pack x: %v\n%s", err, out)
	}
	// Substitute the gasm object for the assembler's archive member (pack
	// extracts members read-only).
	member := filepath.Join(membersDir, asmMember)
	if err := os.Chmod(member, 0o644); err != nil {
		t.Fatal(err)
	}
	if err := os.WriteFile(member, obj, 0o644); err != nil {
		t.Fatal(err)
	}
	listCmd := exec.Command(goBin, "tool", "pack", "t", pkgArch)
	listCmd.Env = append(os.Environ(), "GOARCH=loong64")
	listOut, err := listCmd.CombinedOutput()
	if err != nil {
		t.Fatalf("pack t: %v\n%s", err, listOut)
	}
	newArch := filepath.Join(dir, "pkg.a")
	args := []string{"tool", "pack", "c", newArch}
	seen := map[string]bool{}
	for m := range strings.FieldsSeq(string(listOut)) {
		if seen[m] {
			continue
		}
		seen[m] = true
		if err := os.Chmod(filepath.Join(membersDir, m), 0o644); err != nil {
			t.Fatal(err)
		}
		args = append(args, filepath.Join(membersDir, m))
	}
	pack := exec.Command(goBin, args...)
	pack.Dir = membersDir
	pack.Env = append(os.Environ(), "GOARCH=loong64")
	if out, err := pack.CombinedOutput(); err != nil {
		t.Fatalf("pack c: %v\n%s", err, out)
	}

	// Re-link with our archive in place of the toolchain's.  The link line
	// carries a GOROOT assignment and $WORK placeholders; run it through the
	// shell with the GOEXPERIMENT and GOARCH the toolchain expects (the
	// linker compares the object header against its own, experiments
	// included).
	linkLine = strings.ReplaceAll(linkLine, "$WORK", work)
	linkLine = strings.ReplaceAll(linkLine, filepath.Join(work, "b001", "_pkg_.a"), newArch)
	linkLine = strings.ReplaceAll(linkLine, filepath.Join(work, "b001", "exe", "a.out"), filepath.Join(dir, "app2"))
	link := exec.Command("sh", "-c", linkLine)
	link.Dir = dir
	goExp, _ := exec.Command(goBin, "env", "GOEXPERIMENT").Output()
	link.Env = append(os.Environ(), "GOEXPERIMENT="+strings.TrimSpace(string(goExp)), "GOARCH=loong64")
	if out, err := link.CombinedOutput(); err != nil {
		t.Fatalf("link with gasm object: %v\n%s", err, out)
	}

	// The binary is not executed: there is no LoongArch host or qemu here.
	// The link itself and the symbol table prove cmd/link accepted the gasm
	// object and laid out the function.
	nm := exec.Command(goBin, "tool", "nm", filepath.Join(dir, "app2"))
	nm.Env = append(os.Environ(), "GOARCH=loong64")
	nmOut, err := nm.CombinedOutput()
	if err != nil {
		t.Fatalf("nm gasm-linked binary: %v\n%s", err, nmOut)
	}
	if !strings.Contains(string(nmOut), "main.add") {
		t.Errorf("main.add not found in linked binary:\n%s", nmOut)
	}
}