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

package asm

import (
	"encoding/binary"
	"fmt"
)

// AArch64 ELF64 relocatable object emission.

const (
	emAARCH64 = 183 // EM_AARCH64

	// AArch64 relocation types (the ELF psABI).
	rArm64PrelPgHi21   = 275 // R_AARCH64_ADR_PREL_PG_HI21 (ADRP page)
	rArm64AddAbsLo12NC = 277 // R_AARCH64_ADD_ABS_LO12_NC (ADD page offset)
	rArm64Call26       = 283 // R_AARCH64_CALL26 (BL instruction)
	rArm64Ldst64Lo12NC = 286 // R_AARCH64_LDST64_ABS_LO12_NC (64-bit LDR/STR page offset)
	// R_AARCH64_ABS32 (debug/elf 258): the absolute 32-bit address of a
	// symbol, the R_ADDR shape a 4-byte DATA field carries.  ABS64 (257)
	// lives with the DWARF fixup constants as rAARCH64Abs64.
	rArm64Abs32 = 258
)

// ELFAARCH64Object returns the image as an ELF64 relocatable object file for
// AArch64 (EM_AARCH64, 64-bit, little-endian).  The structure mirrors the
// amd64 and RISC-V ELF emitters: .text, .data, .symtab, .strtab, an
// optional .rela.text and an optional .rela.data.
func (img *Image) ELFAARCH64Object() ([]byte, error) {
	le := binary.LittleEndian

	const (
		secText = 1
		secData = 2
	)

	// Build symbol table.
	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{
		{},
		{name: ".text", info: sttSection, shndx: secText},
		{name: ".data", info: sttSection, shndx: secData},
	}
	syms = append(syms, locals...)
	shInfo := len(syms)
	syms = append(syms, globals...)
	symIdx := map[string]int{}
	for i, s := range syms {
		symIdx[s.name] = i
	}

	// Build relocations.  Each SB reference is an ADRP pair:
	//   ADRP Rd, 0  → R_AARCH64_ADR_PREL_PG_HI21 at the ADRP
	//   ADD         → R_AARCH64_ADD_ABS_LO12_NC at the ADD word
	//   LDR/STR X   → R_AARCH64_LDST64_ABS_LO12_NC at the LDR/STR word
	//   BL          → R_AARCH64_CALL26
	// cmd/link's own conversion emits the HI21 at sectoff and the LO12 at
	// sectoff+4 (cmd/link/internal/arm64/asm.go), so the ADD or load word
	// carries the page-offset relocation, never a second HI21.  The
	// assembler records two RelArm64Addr relocs per ADRP+ADD pair (one per
	// word), so the second of the pair is consumed here.
	// Addends stay raw: ADR_PREL_PG_HI21 and the ABS_LO12_NC forms resolve
	// against S+A, and CALL26 branches take the branch instruction's own
	// place as the PC-relative base, so subtracting the field width (the
	// amd64 R_PCREL convention) would misplace every branch by 4 bytes.
	type elfRela struct {
		off    uint64
		typ    uint32
		sym    int
		addend int64
	}
	var relas []elfRela
	for _, fn := range img.Funcs {
		for i := 0; i < len(fn.Relocs); i++ {
			r := fn.Relocs[i]
			idx, ok := symIdx[r.Name]
			if !ok {
				return nil, fmt.Errorf("relocation references unknown symbol %q", r.Name)
			}
			switch r.Kind {
			case RelArm64Branch:
				relas = append(relas, elfRela{
					off: uint64(fn.Offset + r.Off), typ: rArm64Call26, sym: idx, addend: r.Addend,
				})
			case RelArm64Addr:
				// ADRP+ADD: the pair's second reloc (at Off+4) is the
				// assembler's twin of the same pair; skip it.
				relas = append(relas,
					elfRela{off: uint64(fn.Offset + r.Off), typ: rArm64PrelPgHi21, sym: idx, addend: r.Addend},
					elfRela{off: uint64(fn.Offset + r.Off + 4), typ: rArm64AddAbsLo12NC, sym: idx, addend: r.Addend},
				)
				i++
			case RelArm64LDST64:
				// ADRP+LDR/STR: one assembler reloc covers the pair.
				relas = append(relas,
					elfRela{off: uint64(fn.Offset + r.Off), typ: rArm64PrelPgHi21, sym: idx, addend: r.Addend},
					elfRela{off: uint64(fn.Offset + r.Off + 4), typ: rArm64Ldst64Lo12NC, sym: idx, addend: r.Addend},
				)
			default:
				return nil, fmt.Errorf("relocation kind %v unsupported in ELF emission", r.Kind)
			}
		}
	}

	// The data symbols' symbol-valued DATA fields ("DATA s+0(SB)/8,
	// $other(SB)") become .rela.data entries: an absolute relocation of the
	// DATA line's width at the field's data-section offset, S + A with no
	// PC term.  Widths 4 and 8 have ELF relocation shapes; narrower fields
	// cannot hold an address, so they are refused rather than truncated.
	var dataRelas []elfRela
	for _, d := range img.DataSyms {
		for _, r := range d.Relocs {
			idx, ok := symIdx[r.Name]
			if !ok {
				return nil, fmt.Errorf("data relocation references unknown symbol %q", r.Name)
			}
			var typ uint32
			switch r.Siz {
			case 8:
				typ = rAARCH64Abs64
			case 4:
				typ = rArm64Abs32
			default:
				return nil, fmt.Errorf("DATA %q: a symbol value of width %d has no ELF relocation", d.Name, r.Siz)
			}
			dataRelas = append(dataRelas, elfRela{
				off:    uint64(d.Offset + r.Off),
				sym:    idx,
				typ:    typ,
				addend: r.Addend,
			})
		}
	}

	// Section presence: .rela.text only when there are code relocations,
	// .rela.data only when a DATA line holds a symbol value.
	hasRela := len(relas) > 0
	hasDataRela := len(dataRelas) > 0
	nSections := 6
	if hasRela {
		nSections++
	}
	if hasDataRela {
		nSections++
	}
	secSymtab, secStrtab := 3, 4
	secShstr := nSections - 1

	// 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)
	}
	if hasDataRela {
		stSections.add(".rela.data")
	}
	for _, n := range dwarfSectionNames {
		stSections.add(n)
	}

	// Layout.
	var out []byte
	out = append(out, make([]byte, 64)...)

	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
		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, relaDataOff 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|uint64(r.typ))
			le.PutUint64(b[16:], uint64(r.addend))
			out = append(out, b[:]...)
		}
	}
	if hasDataRela {
		align(8)
		relaDataOff = len(out)
		for _, r := range dataRelas {
			var b [24]byte
			le.PutUint64(b[0:], r.off)
			le.PutUint64(b[8:], uint64(r.sym)<<32|uint64(r.typ))
			le.PutUint64(b[16:], uint64(r.addend))
			out = append(out, b[:]...)
		}
	}

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

	// DWARF debug sections; the address placeholders they leave are carried
	// as .rela.debug_info/.rela.debug_line entries the system linker applies.
	dwAlign := func(n int) {
		for len(out)%n != 0 {
			out = append(out, 0)
		}
	}
	dw := appendDWARFSections(&out, img, dwarfSourceName(img), symIdx, dwAlign, cfiARM64)
	dwarfStart := 0 // section index of .debug_abbrev, set when DWARF is present
	if dw != nil {
		// Five DWARF sections: .debug_abbrev, .debug_info, .debug_line,
		// .debug_line_str and .debug_frame (the CIE is unconditional, so
		// the frame section is always present), plus the relocation
		// sections below when they carry entries.
		dwarfStart = nSections
		nSections += 5
		appendDWARFRelas(&out, dw, rAARCH64Abs64, dwAlign)
		if dw.infoRelaCount > 0 {
			nSections++
		}
		if dw.lineRelaCount > 0 {
			nSections++
		}
		if dw.frameRelaCount > 0 {
			nSections++
		}
	}

	align(8)
	shoff := len(out)

	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)
		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)
	}
	if hasDataRela {
		putSh(".rela.data", shtRela, 0, relaDataOff, 24*len(dataRelas), secSymtab, secData, 8, 24)
	}
	putSh(".shstrtab", shtStrtab, 0, shstrOff, len(stSections.bytes()), 0, 0, 1, 0)
	// DWARF section headers; their indices follow the write order.
	if dw != nil {
		// secIdx is a running section index: each putSh below emits the
		// next header, and the sh_info of a .rela section names the index
		// of the section it relocates.
		secIdx := dwarfStart
		putSh(".debug_abbrev", shtProgbits, 0, dw.abbrevOff, dw.abbrevSize, 0, 0, 1, 0)
		secIdx++
		putSh(".debug_info", shtProgbits, 0, dw.infoOff, dw.infoSize, 0, 0, 1, 0)
		secInfoIdx := secIdx
		secIdx++
		if dw.infoRelaCount > 0 {
			putSh(".rela.debug_info", shtRela, 0, dw.infoRelaOff, 24*dw.infoRelaCount, secSymtab, secInfoIdx, 8, 24)
			secIdx++
		}
		putSh(".debug_line", shtProgbits, 0, dw.lineOff, dw.lineSize, 0, 0, 1, 0)
		secLineIdx := secIdx
		secIdx++
		if dw.lineRelaCount > 0 {
			putSh(".rela.debug_line", shtRela, 0, dw.lineRelaOff, 24*dw.lineRelaCount, secSymtab, secLineIdx, 8, 24)
			secIdx++
		}
		putSh(".debug_line_str", shtProgbits, 0, dw.lineStrOff, dw.lineStrSize, 0, 0, 1, 0)
		secIdx++
		if dw.frameSize > 0 {
			putSh(".debug_frame", shtProgbits, 0, dw.frameOff, dw.frameSize, 0, 0, 8, 0)
			secFrameIdx := secIdx
			secIdx++
			if dw.frameRelaCount > 0 {
				putSh(".rela.debug_frame", shtRela, 0, dw.frameRelaOff, 24*dw.frameRelaCount, secSymtab, secFrameIdx, 8, 24)
			}
		}
	}

	// 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:], emAARCH64)
	le.PutUint32(hdr[20:], elfVersion)
	le.PutUint64(hdr[24:], 0)
	le.PutUint64(hdr[32:], 0)
	le.PutUint64(hdr[40:], uint64(shoff))
	le.PutUint32(hdr[48:], 0)
	le.PutUint16(hdr[52:], 64)
	le.PutUint16(hdr[54:], 0)
	le.PutUint16(hdr[56:], 0)
	le.PutUint16(hdr[58:], 64)
	le.PutUint16(hdr[60:], uint16(nSections))
	le.PutUint16(hdr[62:], uint16(secShstr))

	return out, nil
}