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

package asm

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

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

// TestGOObjectRISCVCallReloc checks that CALL sym(SB) emits a single JAL
// instruction carrying an R_RISCV_JAL relocation (4-byte field) in both the
// GOOBJ and ELF object emitters.
func TestGOObjectRISCVCallReloc(t *testing.T) {
	f, errs := parser.Parse("k_riscv64.s", `
#include "textflag.h"

TEXT ·c(SB), NOSPLIT, $0-0
	CALL callee<>(SB)
	RET

GLOBL callee<>(SB), RODATA, $8
DATA callee<>+0(SB)/8, $42
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileRISCV(f)
	if err != nil {
		t.Fatalf("AssembleFileRISCV: %v", err)
	}
	fn := img.Funcs[0]
	if len(fn.Relocs) != 1 {
		t.Fatalf("relocs = %d, want 1", len(fn.Relocs))
	}
	r := fn.Relocs[0]
	if r.Kind != RelRISCVJal || r.Off != 8 || r.After != 12 || r.Name != "callee" || r.Addend != 0 || r.External {
		t.Errorf("reloc = {kind %v off %d after %d name %q addend %d external %v}", r.Kind, r.Off, r.After, r.Name, r.Addend, r.External)
	}

	obj, err := img.GOObjectRISCV("testpkg", "k_riscv64.s")
	if err != nil {
		t.Fatalf("GOObjectRISCV: %v", err)
	}
	v := openGoobj(t, obj)
	relocIdx := v.blk(blkRelocIdx)
	relocs := v.blk(blkReloc)
	// The function is the last non-package symbol: 4 package defs, then the
	// 4 pc tables and the function.
	first := int(binary.LittleEndian.Uint32(relocIdx[(4+4)*4:]))
	if (first+1)*23 > len(relocs) {
		t.Fatalf("reloc block too short: first=%d len=%d", first, len(relocs))
	}
	e := relocs[first*23:]
	le := binary.LittleEndian
	if int32(le.Uint32(e[0:])) != 8 || e[4] != 4 || le.Uint16(e[5:]) != relocRISCVJal || le.Uint32(e[15:]) != pkgIdxSelf || le.Uint32(e[19:]) != 0 {
		t.Errorf("GOOBJ reloc = off %d size %d type %d pkg %d sym %d", int32(le.Uint32(e[0:])), e[4], le.Uint16(e[5:]), le.Uint32(e[15:]), le.Uint32(e[19:]))
	}

	// The ELF object must carry a single R_RISCV_JAL relocation in .rela.text.
	elfObj, err := img.ELFRISCVObject()
	if err != nil {
		t.Fatalf("ELFRISCVObject: %v", err)
	}
	if !hasELFRISCVJAL(t, elfObj) {
		t.Error("ELF object missing R_RISCV_JAL relocation")
	}
}

// TestELFRISCVPCRELLO12Anchor checks the psABI's LO12 pairing rule: the
// R_RISCV_PCREL_LO12_I/S relocation must reference a symbol whose value is
// the AUIPC site of its HI20 partner (psABI §8.4.9; cmd/link generates one
// local text symbol per AUIPC for exactly this).  The emitter pairs each
// HI20 (against the target symbol) with a LO12 against the .text section
// symbol whose addend is the AUIPC's section-relative offset, so S + A is
// the AUIPC address.
func TestELFRISCVPCRELLO12Anchor(t *testing.T) {
	f, errs := parser.Parse("k_riscv64.s", `
#include "textflag.h"

TEXT ·sb(SB), NOSPLIT, $0-0
	MOV $answer<>(SB), X10
	MOV answer<>(SB), X11
	MOV X12, answer<>(SB)
	RET

GLOBL answer<>(SB), RODATA, $8
DATA answer<>+0(SB)/8, $42
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileRISCV(f)
	if err != nil {
		t.Fatalf("AssembleFileRISCV: %v", err)
	}
	obj, err := img.ELFRISCVObject()
	if err != nil {
		t.Fatalf("ELFRISCVObject: %v", err)
	}
	ef, err := elf.NewFile(bytes.NewReader(obj))
	if err != nil {
		t.Fatalf("parse ELF: %v", err)
	}
	defer ef.Close()
	if flags := binary.LittleEndian.Uint32(obj[48:]); flags != efRISCVFloatAbiDouble {
		t.Errorf("e_flags = %#x, want %#x (EF_RISCV_FLOAT_ABI_DOUBLE)", flags, efRISCVFloatAbiDouble)
	}
	rela := ef.Section(".rela.text")
	if rela == nil {
		t.Fatal("missing .rela.text")
	}
	b, err := rela.Data()
	if err != nil {
		t.Fatal(err)
	}
	if len(b) != 6*24 {
		t.Fatalf(".rela.text holds %d entries, want six (three HI20/LO12 pairs)", len(b)/24)
	}
	le := binary.LittleEndian
	wantLo := []uint32{rRISCVPCRELLO12I, rRISCVPCRELLO12I, rRISCVPCRELLO12S}
	for p := range 3 {
		auipc := 8 * p
		hi := b[p*2*24:]
		lo := b[(p*2+1)*24:]
		if off := le.Uint64(hi[0:]); off != uint64(auipc) {
			t.Errorf("pair %d: HI20 r_offset = %d, want %d (the AUIPC)", p, off, auipc)
		}
		if typ := uint32(le.Uint64(hi[8:])); typ != rRISCVPCRELHI20 {
			t.Errorf("pair %d: HI20 type = %d, want %d", p, typ, rRISCVPCRELHI20)
		}
		if sym := int(le.Uint64(hi[8:]) >> 32); sym == 0 || sym == 1 {
			t.Errorf("pair %d: HI20 against symbol %d, want the target", p, sym)
		}
		if off := le.Uint64(lo[0:]); off != uint64(auipc+4) {
			t.Errorf("pair %d: LO12 r_offset = %d, want %d", p, off, auipc+4)
		}
		if typ := uint32(le.Uint64(lo[8:])); typ != wantLo[p] {
			t.Errorf("pair %d: LO12 type = %d, want %d", p, typ, wantLo[p])
		}
		// The LO12 must denote the AUIPC site: the .text section symbol
		// (index 1) plus the AUIPC's section-relative offset as addend.
		if sym := int(le.Uint64(lo[8:]) >> 32); sym != 1 {
			t.Errorf("pair %d: LO12 against symbol %d, want 1 (the .text section symbol)", p, sym)
		}
		if add := int64(le.Uint64(lo[16:])); add != int64(auipc) {
			t.Errorf("pair %d: LO12 addend = %d, want %d (S + A = the AUIPC address)", p, add, auipc)
		}
	}
}

func TestGOObjectRISCVStructure(t *testing.T) {
	f, errs := parser.Parse("k_riscv64.s", `
#include "textflag.h"

TEXT ·sb(SB), NOSPLIT, $0-0
	MOV $answer<>(SB), X10
	MOV answer<>(SB), X11
	MOV X12, answer<>(SB)
	RET

GLOBL answer<>(SB), RODATA, $8
DATA answer<>+0(SB)/8, $42
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileRISCV(f)
	if err != nil {
		t.Fatalf("AssembleFileRISCV: %v", err)
	}
	fn := img.Funcs[0]
	if fn.Size != 28 {
		t.Fatalf("function size = %d, want 28", fn.Size)
	}
	if len(fn.Relocs) != 3 {
		t.Fatalf("relocs = %d, want 3", len(fn.Relocs))
	}
	wantKind := []RelocKind{RelRISCVPCRELIType, RelRISCVPCRELIType, RelRISCVPCRELSType}
	wantOff := []int{0, 8, 16}
	for i, r := range fn.Relocs {
		if r.Kind != wantKind[i] || r.Off != wantOff[i] || r.After != r.Off+8 || r.Name != "answer" || r.Addend != 0 {
			t.Errorf("reloc %d = {kind %v off %d after %d name %q addend %d}", i, r.Kind, r.Off, r.After, r.Name, r.Addend)
		}
	}

	obj, err := img.GOObjectRISCV("testpkg", "k_riscv64.s")
	if err != nil {
		t.Fatalf("GOObjectRISCV: %v", err)
	}
	v := openGoobj(t, obj)

	// Package defs: the static GLOBL, the FuncInfo, then the two DWARF
	// symbols.
	defs := v.syms(blkSymdef)
	if len(defs) != 4 {
		t.Fatalf("symdefs = %d, want 4", len(defs))
	}
	if defs[0].name != "answer" || defs[0].abi != 0xffff || defs[0].typ != kindSRODATA || defs[0].size != 8 {
		t.Errorf("answer symbol = %+v", defs[0])
	}
	if defs[2].typ != kindSDWARFLINES || defs[3].typ != kindSDWARFFCN {
		t.Errorf("dwarf symbols = %+v, %+v", defs[2], defs[3])
	}

	// The three code relocations, in definition order: ITYPE, ITYPE, STYPE,
	// each 8 bytes wide against the GLOBL (package symbol 0).
	relocIdx := v.blk(blkRelocIdx)
	relocs := v.blk(blkReloc)
	if len(relocs) != 5*23 {
		t.Fatalf("relocs = %d bytes, want 5 entries", len(relocs))
	}
	// The function is the last non-package symbol; its relocs start after
	// the DWARF symbols' (defs 2 and 3 each carry one).
	le := binary.LittleEndian
	first := int(le.Uint32(relocIdx[4*(4+4):]))
	wantType := []uint16{relocRISCVPcrelItype, relocRISCVPcrelItype, relocRISCVPcrelStype}
	wantOffAbs := []int{0, 8, 16}
	for i := range 3 {
		e := relocs[(first+i)*23:]
		if int32(le.Uint32(e[0:])) != int32(wantOffAbs[i]) || e[4] != 8 || le.Uint16(e[5:]) != wantType[i] ||
			le.Uint32(e[15:]) != pkgIdxSelf || le.Uint32(e[19:]) != 0 {
			t.Errorf("reloc %d = off %d size %d type %d pkg %d sym %d", i, int32(le.Uint32(e[0:])), e[4], le.Uint16(e[5:]), le.Uint32(e[15:]), le.Uint32(e[19:]))
		}
	}

	// The function code: three AUIPC+second-instruction pairs with zero
	// immediates, then the uncompressed JALR X0, 0(X1) the toolchain emits
	// for RET.
	code := img.Code[fn.Offset : fn.Offset+fn.Size]
	want := append(wordLE(riscvUType(riscvEnc{0x17, 0x0, 0x00}, 10, 0)), wordLE(riscvIType(riscvEnc{0x13, 0x0, 0x00}, 10, 10, 0))...)
	want = append(want, wordLE(riscvUType(riscvEnc{0x17, 0x0, 0x00}, 11, 0))...)
	want = append(want, wordLE(riscvIType(riscvEnc{0x03, 0x3, 0x00}, 11, 11, 0))...)
	want = append(want, wordLE(riscvUType(riscvEnc{0x17, 0x0, 0x00}, 31, 0))...)
	want = append(want, wordLE(riscvSType(riscvEnc{0x23, 0x3, 0x00}, 31, 12, 0))...)
	want = append(want, 0x67, 0x80, 0x00, 0x00) // JALR X0, 0(X1)
	if !bytes.Equal(code, want) {
		t.Errorf("code = % x\nwant % x", code, want)
	}

	// The same bytes must survive into the object's data block intact: the
	// linker patches only the immediate fields of the AUIPC pairs, so the
	// opcode/register bits of every instruction must not be zeroed.
	dataIdx := v.blk(blkDataIdx)
	dataBlk := v.blk(blkData)
	dOff := int(le.Uint32(dataIdx[8*4:])) // the function is the last symbol
	emitted := dataBlk[dOff : dOff+fn.Size]
	if !bytes.Equal(emitted, want) {
		t.Errorf("emitted data = % x\nwant % x", emitted, want)
	}
}

// TestGOObjectRISCVLink cross-compiles a Go program with the gasm-produced
// object substituted into the package archive, proving cmd/link accepts the
// emitted RISC-V GOOBJ.  The binary is not executed (no riscv64 host or
// qemu).  Skipped when no Go toolchain is available.
func TestGOObjectRISCVLink(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
	MOV a+0(FP), X10
	MOV b+8(FP), X11
	ADD X11, X10, X10
	MOV X10, ret+16(FP)
	RET
`
	if err := os.WriteFile(filepath.Join(dir, "main_riscv64.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 rvlink\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=riscv64")
	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_riscv64.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)
	asmMember := filepath.Base(strings.ReplaceAll(asmObj, "$WORK", work))

	pf, perrs := parser.Parse(filepath.Join(dir, "main_riscv64.s"), asmSrc)
	if len(perrs) > 0 {
		t.Fatalf("parse: %v", perrs)
	}
	pimg, err := AssembleFileRISCV(pf)
	if err != nil {
		t.Fatalf("AssembleFileRISCV: %v", err)
	}
	obj, err := pimg.GOObjectRISCV("main", filepath.Join(dir, "main_riscv64.s"))
	if err != nil {
		t.Fatalf("GOObjectRISCV: %v", err)
	}

	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=riscv64")
	if out, err := extract.CombinedOutput(); err != nil {
		t.Fatalf("pack x: %v\n%s", err, out)
	}
	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=riscv64")
	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=riscv64")
	if out, err := pack.CombinedOutput(); err != nil {
		t.Fatalf("pack c: %v\n%s", err, out)
	}

	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=riscv64")
	if out, err := link.CombinedOutput(); err != nil {
		t.Fatalf("link with gasm object: %v\n%s", err, out)
	}

	nm := exec.Command(goBin, "tool", "nm", filepath.Join(dir, "app2"))
	nm.Env = append(os.Environ(), "GOARCH=riscv64")
	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)
	}
}

// hasELFRISCVJAL reports whether the ELF object carries an R_RISCV_JAL
// relocation in its .rela.text section.
func hasELFRISCVJAL(t *testing.T, data []byte) bool {
	t.Helper()
	f, err := elf.NewFile(bytes.NewReader(data))
	if err != nil {
		t.Fatalf("parse ELF: %v", err)
	}
	defer f.Close()
	rela := f.Section(".rela.text")
	if rela == nil {
		return false
	}
	b, err := rela.Data()
	if err != nil {
		t.Fatalf(".rela.text data: %v", err)
	}
	const rRISCVJAL = 17
	for i := 0; i+24 <= len(b); i += 24 {
		info := binary.LittleEndian.Uint64(b[i+8:])
		if uint32(info) == rRISCVJAL {
			return true
		}
	}
	return false
}