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

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

// The object-file tests share one source: two exported functions, one
// file-local constant reached through a relocation, and one external symbol
// the linker must resolve.  The functions take their arguments in the System
// V registers (not the Go stack ABI) so a C driver can call them directly.
const elfTestSrc = `
#include "textflag.h"

TEXT ·addq(SB), NOSPLIT, $0
	LEAQ (DI)(SI*1), AX
	RET

TEXT ·getanswer(SB), NOSPLIT, $0
	MOVQ answer<>(SB), AX
	RET

TEXT ·useextern(SB), NOSPLIT, $0
	MOVQ extvar(SB), AX
	RET

GLOBL answer<>(SB), RODATA, $8
DATA answer<>+0(SB)/8, $42
`

func elfTestImage(t *testing.T) *Image {
	t.Helper()
	f, errs := parser.Parse("t_amd64.s", elfTestSrc)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("AssembleFile: %v", err)
	}
	return img
}

// TestAssembleFileExternals checks that a reference to a symbol no GLOBL
// defines is recorded as an external relocation instead of failing — the
// raw image leaves the displacement zero, the object emitters carry it.
func TestAssembleFileExternals(t *testing.T) {
	img := elfTestImage(t)
	if len(img.Externals) != 1 || img.Externals[0] != "extvar" {
		t.Fatalf("Externals = %v, want [extvar]", img.Externals)
	}
	var ext, local int
	for _, fn := range img.Funcs {
		for _, r := range fn.Relocs {
			if r.External {
				ext++
				if r.Name != "extvar" {
					t.Errorf("external reloc names %q, want extvar", r.Name)
				}
			} else {
				local++
				if r.Name != "answer" {
					t.Errorf("local reloc names %q, want answer", r.Name)
				}
			}
		}
	}
	if ext != 1 || local != 1 {
		t.Errorf("relocs = %d external, %d local; want 1 and 1", ext, local)
	}
}

// TestELFObject checks the structure of the emitted ELF64 relocatable
// object: sections, the symbol table (bindings, types, values, sizes) and
// the .rela.text relocations, parsed back with debug/elf.
func TestELFObject(t *testing.T) {
	img := elfTestImage(t)
	obj, err := img.ELFObject()
	if err != nil {
		t.Fatalf("ELFObject: %v", err)
	}
	f, err := elf.NewFile(bytes.NewReader(obj))
	if err != nil {
		t.Fatalf("parse emitted object: %v", err)
	}
	defer f.Close()

	if f.Type != elf.ET_REL || f.Machine != elf.EM_X86_64 {
		t.Errorf("type/machine = %v/%v, want ET_REL/EM_X86_64", f.Type, f.Machine)
	}

	text := f.Section(".text")
	data := f.Section(".data")
	if text == nil || data == nil {
		t.Fatal("missing .text or .data section")
	}
	if text.Flags&elf.SHF_EXECINSTR == 0 || text.Flags&elf.SHF_ALLOC == 0 {
		t.Errorf(".text flags = %v", text.Flags)
	}
	if data.Flags&elf.SHF_WRITE == 0 {
		t.Errorf(".data flags = %v", data.Flags)
	}
	textData, err := text.Data()
	if err != nil {
		t.Fatal(err)
	}
	if !bytes.Equal(textData, img.Code) {
		t.Errorf(".text contents differ from the image code")
	}

	syms, err := f.Symbols()
	if err != nil {
		t.Fatalf("symbols: %v", err)
	}
	byName := map[string]elf.Symbol{}
	for _, s := range syms {
		byName[s.Name] = s
	}
	wantSym := func(name string, bind elf.SymBind, typ elf.SymType, section elf.SectionIndex, size uint64) {
		t.Helper()
		s, ok := byName[name]
		if !ok {
			t.Errorf("symbol %q not found", name)
			return
		}
		if elf.ST_BIND(s.Info) != bind || elf.ST_TYPE(s.Info) != typ {
			t.Errorf("%s: bind/type = %v/%v, want %v/%v", name, elf.ST_BIND(s.Info), elf.ST_TYPE(s.Info), bind, typ)
		}
		if s.Section != section {
			t.Errorf("%s: section = %v, want %v", name, s.Section, section)
		}
		if s.Size != size {
			t.Errorf("%s: size = %d, want %d", name, s.Size, size)
		}
	}
	// The emitted layout is fixed: 0 NULL, 1 .text, 2 .data.
	if f.Sections[1].Name != ".text" || f.Sections[2].Name != ".data" {
		t.Fatalf("section layout = %s, %s; want .text, .data", f.Sections[1].Name, f.Sections[2].Name)
	}
	textIdx := elf.SectionIndex(1)
	dataIdx := elf.SectionIndex(2)
	wantSym("addq", elf.STB_GLOBAL, elf.STT_FUNC, textIdx, 5)
	wantSym("getanswer", elf.STB_GLOBAL, elf.STT_FUNC, textIdx, 8)
	wantSym("useextern", elf.STB_GLOBAL, elf.STT_FUNC, textIdx, 8)
	wantSym("answer", elf.STB_LOCAL, elf.STT_OBJECT, dataIdx, 8)
	wantSym("extvar", elf.STB_GLOBAL, elf.STT_NOTYPE, elf.SHN_UNDEF, 0)

	// Relocations: one for the file-local constant (resolving against the
	// local data symbol) and one for the external (against the undefined
	// global), both R_X86_64_PC32 with the −4 addend the PC-relative form
	// needs.  debug/elf does not surface rela entries, so read the section
	// directly.
	relaSec := f.Section(".rela.text")
	if relaSec == nil {
		t.Fatal("missing .rela.text")
	}
	raw, err := relaSec.Data()
	if err != nil {
		t.Fatal(err)
	}
	if len(raw)%24 != 0 || len(raw)/24 != 2 {
		t.Fatalf(".rela.text has %d bytes, want two 24-byte entries", len(raw))
	}
	// Symbol names straight from the raw tables: r_info carries an index
	// into .symtab including the null entry, which debug/elf's Symbols()
	// slice may not mirror.
	symtabRaw, err := f.Section(".symtab").Data()
	if err != nil {
		t.Fatal(err)
	}
	strtabRaw, err := f.Section(".strtab").Data()
	if err != nil {
		t.Fatal(err)
	}
	symName := func(idx int) string {
		stName := binary.LittleEndian.Uint32(symtabRaw[idx*24:])
		end := bytes.IndexByte(strtabRaw[stName:], 0)
		return string(strtabRaw[stName : int(stName)+end])
	}
	for i := 0; i < 2; i++ {
		e := raw[i*24 : (i+1)*24]
		off := binary.LittleEndian.Uint64(e[0:])
		info := binary.LittleEndian.Uint64(e[8:])
		addend := int64(binary.LittleEndian.Uint64(e[16:]))
		typ := info & 0xffffffff
		sym := int(info >> 32)
		if typ != uint64(elf.R_X86_64_PC32) {
			t.Errorf("reloc %d: type %d, want R_X86_64_PC32", i, typ)
		}
		if addend != -4 {
			t.Errorf("reloc %d: addend %d, want -4", i, addend)
		}
		if name := symName(sym); name != "answer" && name != "extvar" {
			t.Errorf("reloc %d: symbol %q, want answer or extvar", i, name)
		}
		// The relocation offset lands on the disp32 field: the four bytes
		// before a RET-terminated eight-byte MOVQ.
		if off+4 > uint64(len(textData)) {
			t.Errorf("reloc %d: offset %d outside .text", i, off)
		}
	}
}

// TestELFObjectNoRelocations checks a file with no static-symbol references
// emits a valid object without a .rela.text section.
func TestELFObjectNoRelocations(t *testing.T) {
	f, errs := parser.Parse("n_amd64.s", `
#include "textflag.h"
TEXT ·nop(SB), NOSPLIT, $0
	RET
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("AssembleFile: %v", err)
	}
	obj, err := img.ELFObject()
	if err != nil {
		t.Fatalf("ELFObject: %v", err)
	}
	ef, err := elf.NewFile(bytes.NewReader(obj))
	if err != nil {
		t.Fatalf("parse emitted object: %v", err)
	}
	defer ef.Close()
	if ef.Section(".rela.text") != nil {
		t.Error("unexpected .rela.text section")
	}
	syms, err := ef.Symbols()
	if err != nil {
		t.Fatal(err)
	}
	found := false
	for _, s := range syms {
		if s.Name == "nop" && elf.ST_TYPE(s.Info) == elf.STT_FUNC {
			found = true
		}
	}
	if !found {
		t.Error("function symbol nop not found")
	}
}

// TestELFLinkAndRun is the end-to-end check: assemble the test functions,
// link the emitted object with a C driver that defines the external symbol,
// and run the result.  Skipped when no C compiler is available.
func TestELFLinkAndRun(t *testing.T) {
	cc, err := exec.LookPath("cc")
	if err != nil {
		t.Skip("no C compiler available")
	}
	dir := t.TempDir()

	img := elfTestImage(t)
	obj, err := img.ELFObject()
	if err != nil {
		t.Fatalf("ELFObject: %v", err)
	}
	objPath := filepath.Join(dir, "t.o")
	if err := os.WriteFile(objPath, obj, 0o644); err != nil {
		t.Fatal(err)
	}

	const driver = `
#include <stdio.h>

long addq(long a, long b);
long getanswer(void);
long useextern(void);

long extvar = 7;

int main(void) {
	printf("%ld %ld %ld\n", addq(41, 1), getanswer(), useextern());
	return 0;
}
`
	driverPath := filepath.Join(dir, "driver.c")
	if err := os.WriteFile(driverPath, []byte(driver), 0o644); err != nil {
		t.Fatal(err)
	}

	// -no-pie: the encoder emits R_X86_64_PC32 for external references,
	// which a position-independent executable would reject (it wants
	// PLT32/GOT relocations, a future increment).
	appPath := filepath.Join(dir, "app")
	out, err := exec.Command(cc, "-no-pie", "-o", appPath, driverPath, objPath).CombinedOutput()
	if err != nil {
		t.Fatalf("link failed: %v\n%s", err, out)
	}
	run, err := exec.Command(appPath).CombinedOutput()
	if err != nil {
		t.Fatalf("run failed: %v\n%s", err, run)
	}
	if got := string(run); got != "42 42 7\n" {
		t.Errorf("output %q, want \"42 42 7\\n\"", got)
	}
}