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

package asm

import (
	"encoding/binary"
	"fmt"
)

// This file emits ELF64 relocatable objects (ET_REL) from an assembled
// Image: a .text section holding the function bodies, a .data section
// holding the GLOBL initialisers, a symbol table with one symbol per TEXT
// and GLOBL (file-local <> symbols are STB_LOCAL, the rest STB_GLOBAL), and
// a .rela.text relocation table — one R_X86_64_PC32 entry per static-symbol
// reference, internal references resolving against the local data symbols
// and external ones against undefined globals.  The output links with the
// system toolchain (cc/ld) the way a hand-assembled .o would.

const (
	elfClass64 = 2
	elfDataLSB = 1
	elfVersion = 1

	etREL   = 1 // relocatable object
	emX8664 = 62

	shtNull     = 0
	shtProgbits = 1
	shtSymtab   = 2
	shtStrtab   = 3
	shtRela     = 4

	shfWrite     = 1
	shfAlloc     = 2
	shfExecInstr = 4

	stbGlobal = 1

	sttObject   = 1
	sttFunc     = 2
	sttSection  = 3
	stInfoShift = 4

	rX8664PC32 = 2
)

// elfSym is one symbol-table entry in construction.
type elfSym struct {
	name  string
	info  byte
	shndx uint16
	value uint64
	size  uint64
}

// ELFObject returns the image as an ELF64 relocatable object file, ready for
// the system linker.  Symbol names are the TEXT and GLOBL identifiers as
// written (the middle dot stripped); a package prefix, when present, is
// joined with a dot.  Every static-symbol reference becomes an
// R_X86_64_PC32 relocation, so the code is position-independent and links
// at any address.
func (img *Image) ELFObject() ([]byte, error) {
	le := binary.LittleEndian

	// Section indices: 0 NULL, 1 .text, 2 .data; the tables follow.
	const (
		secText = 1
		secData = 2
	)

	// Build the symbol table: the null entry and the two section symbols
	// come first, then the local symbols (static TEXT and GLOBL), then the
	// globals (exported TEXT and GLOBL, and the undefined externals) — ELF
	// requires every local to precede every global, and sh_info records the
	// boundary.  symIdx maps a symbol name to its index for the relocations.
	var locals, globals []elfSym
	for _, fn := range img.Funcs {
		s := elfSym{
			name:  objectName(fn.Pkg, fn.Name),
			info:  sttFunc,
			shndx: secText,
			value: uint64(fn.Offset),
			size:  uint64(fn.Size),
		}
		if fn.Static {
			locals = append(locals, s)
		} else {
			s.info |= stbGlobal << stInfoShift
			globals = append(globals, s)
		}
	}
	for _, d := range img.DataSyms {
		s := elfSym{
			name:  objectName(d.Pkg, d.Name),
			info:  sttObject,
			shndx: secData,
			value: uint64(d.Offset),
			size:  uint64(d.Size),
		}
		if d.Static {
			locals = append(locals, s)
		} else {
			s.info |= stbGlobal << stInfoShift
			globals = append(globals, s)
		}
	}
	for _, name := range img.Externals {
		globals = append(globals, elfSym{name: name, info: stbGlobal << stInfoShift})
	}
	syms := []elfSym{
		{}, // the mandatory null entry
		{name: ".text", info: sttSection, shndx: secText},
		{name: ".data", info: sttSection, shndx: secData},
	}
	syms = append(syms, locals...)
	shInfo := len(syms) // first global symbol
	syms = append(syms, globals...)
	symIdx := map[string]int{}
	for i, s := range syms {
		symIdx[s.name] = i
	}

	// Build the relocations.
	type elfRela struct {
		off    uint64
		sym    int
		addend int64
	}
	var relas []elfRela
	for _, fn := range img.Funcs {
		for _, r := range fn.Relocs {
			idx, ok := symIdx[r.Name]
			if !ok {
				return nil, fmt.Errorf("relocation references unknown symbol %q", r.Name)
			}
			relas = append(relas, elfRela{
				off: uint64(fn.Offset + r.Off),
				sym: idx,
				// R_X86_64_PC32 computes S + A − P with P the patch site; the
				// assembler measures the symbol from the instruction end,
				// After − Off bytes past the field, so the addend carries
				// that distance with a negative sign.
				addend: r.Addend - int64(r.After-r.Off),
			})
		}
	}

	// Serialise the string tables.
	stNames := newElfStrtab()
	for _, s := range syms {
		stNames.add(s.name)
	}
	stSections := newElfStrtab()
	for _, n := range []string{".text", ".data", ".symtab", ".strtab", ".rela.text", ".shstrtab"} {
		stSections.add(n)
	}
	for _, n := range dwarfSectionNames {
		stSections.add(n)
	}

	// Section presence: .rela.text only when there are relocations.
	hasRela := len(relas) > 0
	nSections := 6 // NULL, .text, .data, .symtab, .strtab, .shstrtab
	if hasRela {
		nSections = 7
	}
	secSymtab, secStrtab := 3, 4
	secShstr := nSections - 1

	// Lay the file out: header, section data, section headers.
	var out []byte
	out = append(out, make([]byte, 64)...) // ELF header, filled last

	align := func(n int) {
		for len(out)%n != 0 {
			out = append(out, 0)
		}
	}

	align(16)
	textOff := len(out)
	out = append(out, img.Code...)

	align(16)
	dataOff := len(out)
	out = append(out, img.Data...)

	align(8)
	symtabOff := len(out)
	for _, s := range syms {
		var b [24]byte
		le.PutUint32(b[0:], uint32(stNames.at(s.name)))
		b[4] = s.info
		b[5] = 0 // st_other
		le.PutUint16(b[6:], s.shndx)
		le.PutUint64(b[8:], s.value)
		le.PutUint64(b[16:], s.size)
		out = append(out, b[:]...)
	}

	strtabOff := len(out)
	out = append(out, stNames.bytes()...)

	var relaOff int
	if hasRela {
		align(8)
		relaOff = len(out)
		for _, r := range relas {
			var b [24]byte
			le.PutUint64(b[0:], r.off)
			le.PutUint64(b[8:], uint64(r.sym)<<32|rX8664PC32)
			le.PutUint64(b[16:], uint64(r.addend))
			out = append(out, b[:]...)
		}
	}

	shstrOff := len(out)
	out = append(out, stSections.bytes()...)

	// DWARF debug sections (no relocations — the linker resolves DWARF fixups).
	dwAlign := func(n int) {
		for len(out)%n != 0 {
			out = append(out, 0)
		}
	}
	dw := appendDWARFSections(&out, img, "gasm.s", symIdx, dwAlign)
	if dw != nil {
		nSections += 4 // .debug_abbrev, .debug_info, .debug_line, .debug_line_str
	}

	align(8)
	shoff := len(out)

	// Section headers.
	putSh := func(name string, typ int, flags uint64, off, size int, link, info int, alignV, entsize uint64) {
		var b [64]byte
		le.PutUint32(b[0:], uint32(stSections.at(name)))
		le.PutUint32(b[4:], uint32(typ))
		le.PutUint64(b[8:], flags)
		le.PutUint64(b[16:], 0) // sh_addr
		le.PutUint64(b[24:], uint64(off))
		le.PutUint64(b[32:], uint64(size))
		le.PutUint32(b[40:], uint32(link))
		le.PutUint32(b[44:], uint32(info))
		le.PutUint64(b[48:], alignV)
		le.PutUint64(b[56:], entsize)
		out = append(out, b[:]...)
	}
	putSh("", shtNull, 0, 0, 0, 0, 0, 0, 0)
	putSh(".text", shtProgbits, shfAlloc|shfExecInstr, textOff, len(img.Code), 0, 0, 16, 0)
	putSh(".data", shtProgbits, shfAlloc|shfWrite, dataOff, len(img.Data), 0, 0, 16, 0)
	putSh(".symtab", shtSymtab, 0, symtabOff, 24*len(syms), secStrtab, shInfo, 8, 24)
	putSh(".strtab", shtStrtab, 0, strtabOff, len(stNames.bytes()), 0, 0, 1, 0)
	if hasRela {
		putSh(".rela.text", shtRela, 0, relaOff, 24*len(relas), secSymtab, secText, 8, 24)
	}
	putSh(".shstrtab", shtStrtab, 0, shstrOff, len(stSections.bytes()), 0, 0, 1, 0)

	// DWARF section headers.
	if dw != nil {
		putSh(".debug_abbrev", shtProgbits, 0, dw.abbrevOff, dw.abbrevSize, 0, 0, 1, 0)
		putSh(".debug_info", shtProgbits, 0, dw.infoOff, dw.infoSize, 0, 0, 1, 0)
		putSh(".debug_line", shtProgbits, 0, dw.lineOff, dw.lineSize, 0, 0, 1, 0)
		putSh(".debug_line_str", shtProgbits, 0, dw.lineStrOff, dw.lineStrSize, 0, 0, 1, 0)
		if dw.frameSize > 0 {
			putSh(".debug_frame", shtProgbits, 0, dw.frameOff, dw.frameSize, 0, 0, 8, 0)
		}
	}

	// The ELF header.
	hdr := out[:64]
	copy(hdr[0:], []byte{0x7f, 'E', 'L', 'F', elfClass64, elfDataLSB, elfVersion, 0})
	le.PutUint16(hdr[16:], etREL)
	le.PutUint16(hdr[18:], emX8664)
	le.PutUint32(hdr[20:], elfVersion)
	le.PutUint64(hdr[24:], 0)             // e_entry
	le.PutUint64(hdr[32:], 0)             // e_phoff
	le.PutUint64(hdr[40:], uint64(shoff)) // e_shoff
	le.PutUint32(hdr[48:], 0)             // e_flags
	le.PutUint16(hdr[52:], 64)            // e_ehsize
	le.PutUint16(hdr[54:], 0)             // e_phentsize
	le.PutUint16(hdr[56:], 0)             // e_phnum
	le.PutUint16(hdr[58:], 64)            // e_shentsize
	le.PutUint16(hdr[60:], uint16(nSections))
	le.PutUint16(hdr[62:], uint16(secShstr))

	return out, nil
}

// objectName renders a symbol's object-file name: the identifier as written,
// with an explicit package prefix joined by a dot.
func objectName(pkg, name string) string {
	if pkg == "" {
		return name
	}
	return pkg + "." + name
}

// elfStrtab is an ELF string table under construction.
type elfStrtab struct {
	buf []byte
	off map[string]int
}

func newElfStrtab() *elfStrtab {
	return &elfStrtab{buf: []byte{0}, off: map[string]int{"": 0}}
}

func (s *elfStrtab) add(name string) {
	if _, ok := s.off[name]; ok {
		return
	}
	s.off[name] = len(s.buf)
	s.buf = append(s.buf, name...)
	s.buf = append(s.buf, 0)
}

func (s *elfStrtab) at(name string) int { return s.off[name] }

func (s *elfStrtab) bytes() []byte { return s.buf }