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

package asm

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"sync"
)

// This file emits GOOBJ — the Go toolchain's object format, which cmd/link
// consumes directly — so gasm-assembled functions drop into a go build
// without the Go assembler.  The layout follows cmd/internal/goobj: a
// toolchain preamble ("go object ...\n!\n"), the go120ld header with its
// block offsets, a string table, symbol definitions, the relocation /
// aux / data index arrays, and the three blocks themselves.
//
// The object carries what the linker requires of an assembly object: the
// functions (non-package symbols, as cmd/asm emits them), the GLOBL data,
// one FuncInfo per function, the per-function DWARF symbols (the
// .debug_line program and the subprogram DIE, which the linker's DWARF
// pass reads verbatim), and the pc-value tables (pcsp, pcfile, pcline,
// pcinline).  The implicit funcdata symbols are omitted; the linker fills
// their defaults.
//
// emitGOObject is architecture-agnostic; the per-architecture GOObject*
// methods supply the toolchain preamble, the MinLC (pc-value delta unit)
// and the relocation-type mapping for code relocations.

// GOOBJ block indices (cmd/internal/goobj).
const (
	blkAutolib = iota
	blkPkgIdx
	blkFile
	blkSymdef
	blkHashed64def
	blkHasheddef
	blkNonpkgdef
	blkNonpkgref
	blkRefFlags
	blkHash64
	blkHash
	blkRelocIdx
	blkAuxIdx
	blkDataIdx
	blkReloc
	blkAux
	blkData
	blkRefName
	blkEnd
)

// Symbol kinds used by assembly objects (cmd/internal/objabi).
const (
	kindSTEXT       = 1
	kindSRODATA     = 3
	kindSDATA       = 7
	kindSDWARFFCN   = 14
	kindSDWARFLINES = 20
)

// Symbol flags (cmd/internal/goobj).
const (
	symFlagDupok   = 0x01
	symFlagNoSplit = 0x10
	symFlag2Link   = 0x10 // asm objects flag every named symbol as linkname
	symABIStatic   = 0xffff
)

// Aux entry types (cmd/internal/goobj).
const (
	auxFuncInfo   = 1
	auxDwarfInfo  = 3
	auxDwarfLines = 6
	auxPcsp       = 7
	auxPcfile     = 8
	auxPcline     = 9
	auxPcinline   = 10
)

// FuncInfo flags (internal/abi).
const (
	funcFlagSPWrite = 2
	funcFlagAsm     = 4
)

// Relocation types (cmd/internal/objabi).
const (
	relocPCRel       = 14  // R_PCREL
	relocAddr        = 1   // R_ADDR
	relocDWTXTADDRU4 = 106 // R_DWTXTADDR_U4
)

// Special package indices for symbol references.
const (
	pkgIdxNone = 0x7fffffff
	pkgIdxSelf = 0x7ffffffb
)

const goobjMagic = "\x00go120ld"

// goSym is one symbol definition under construction.
type goSym struct {
	name  string
	abi   uint16
	typ   uint8
	flag  uint8
	flag2 uint8
	size  uint32
	align uint32
}

func (s goSym) append(b []byte, strOff map[string]uint32) []byte {
	b = binary.LittleEndian.AppendUint32(b, uint32(len(s.name)))
	b = binary.LittleEndian.AppendUint32(b, strOff[s.name])
	b = binary.LittleEndian.AppendUint16(b, s.abi)
	b = append(b, s.typ, s.flag, s.flag2)
	b = binary.LittleEndian.AppendUint32(b, s.size)
	return binary.LittleEndian.AppendUint32(b, s.align)
}

// dwarfRelocSet attaches emitter-generated relocations (the DWARF
// lines/info symbols' address references) to a definition index.
type dwarfRelocSet struct {
	si     int
	relocs []goobjReloc
}

// GOObject returns the image as a GOOBJ object file for the given package
// path (the linker qualifies the exported symbols with it, the way cmd/asm
// does with its -p flag).  srcPath names the source file recorded in the
// object's file table and line tables.  The toolchain's object preamble is
// captured from the installed go tool asm, so the output links with the
// toolchain it was produced on — exactly like a real assembly object.
func (img *Image) GOObject(pkgPath, srcPath string) ([]byte, error) {
	pre, err := toolchainObjectPreamble()
	if err != nil {
		return nil, err
	}
	// amd64: MinLC 1, R_PCREL for the code relocations.
	return img.emitGOObject(pkgPath, srcPath, pre, 1, func(Reloc) (uint16, uint8) { return relocPCRel, 4 })
}

// emitGOObject assembles the GOOBJ payload for any architecture.  pre is
// the toolchain's object preamble; minLC is the architecture's minimum
// instruction length, the unit of the pc-value table deltas; relocField
// maps a code relocation to its objabi relocation type and the width of
// the instruction field the linker writes.
func (img *Image) emitGOObject(pkgPath, srcPath string, pre []byte, minLC int, relocField func(Reloc) (uint16, uint8)) ([]byte, error) {
	if pkgPath == "" {
		return nil, fmt.Errorf("GOOBJ emission requires a package path (-p)")
	}

	// The non-package definitions first — the DWARF symbols reference the
	// functions by these indices: per function the four pc-value tables
	// and the function itself, as cmd/asm lays them out.
	type npSym struct {
		sym  goSym
		data []byte
	}
	var nps []npSym
	type pcRefs struct{ sp, file, line, inl int }
	pcIdx := make([]pcRefs, len(img.Funcs))
	fnNpIdx := make([]int, len(img.Funcs))
	for i, fn := range img.Funcs {
		tables := []struct {
			data []byte
			dst  *int
		}{
			{pcspTable(fn, minLC), &pcIdx[i].sp},
			{pcValueFlat(0, fn.Size, minLC), &pcIdx[i].file},
			{pcValueFlat(int32(fn.Line), fn.Size, minLC), &pcIdx[i].line},
			{pcValueFlat(-1, fn.Size, minLC), &pcIdx[i].inl},
		}
		for _, t := range tables {
			*t.dst = len(nps)
			nps = append(nps, npSym{
				sym:  goSym{typ: kindSRODATA, size: uint32(len(t.data)), align: 1},
				data: t.data,
			})
		}
		name := fn.Name
		abi := uint16(0)
		if fn.Static {
			abi = symABIStatic
		} else {
			name = pkgPath + "." + name
		}
		flag := uint8(0)
		if fn.NoSplit {
			flag |= symFlagNoSplit
		}
		fnNpIdx[i] = len(nps)
		code := append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...)
		for _, r := range fn.Relocs {
			// Only the amd64 encoder resolves file-local static symbols
			// into a disp32 field at assemble time; GOOBJ must leave that
			// field zero for the linker to fill.  The RISC-V and LoongArch
			// encoders emit zero immediates with a relocation instead, and
			// their relocations cover whole AUIPC/pcalau12i pairs, so
			// zeroing r.Off would erase the opcode/register bits the linker
			// preserves when it patches only the immediate.
			if r.Kind != RelPCRel32 {
				continue
			}
			if r.Off >= 0 && r.Off+4 <= len(code) {
				code[r.Off], code[r.Off+1], code[r.Off+2], code[r.Off+3] = 0, 0, 0, 0
			}
		}
		nps = append(nps, npSym{
			sym:  goSym{name: name, abi: abi, typ: kindSTEXT, flag: flag, flag2: symFlag2Link, size: uint32(fn.Size)},
			data: code,
		})
	}

	// The package definitions: the GLOBL symbols, then, per function, the
	// FuncInfo and the two DWARF symbols (the .debug_line program and the
	// subprogram DIE).  defIdx maps a GLOBL's bare name to its definition
	// index for the code relocations.
	var defs []goSym
	var defData [][]byte
	defIdx := map[string]int{}
	for _, d := range img.DataSyms {
		name := d.Name
		if !d.Static {
			name = pkgPath + "." + name
		}
		typ := uint8(kindSDATA)
		if d.Rodata {
			typ = kindSRODATA
		}
		flag := uint8(0)
		if d.Dupok {
			flag = symFlagDupok
		}
		abi := uint16(0)
		if d.Static {
			abi = symABIStatic
		}
		defIdx[d.Name] = len(defs)
		defs = append(defs, goSym{name: name, abi: abi, typ: typ, flag: flag, flag2: symFlag2Link, size: uint32(d.Size)})
		defData = append(defData, img.Data[d.Offset:d.Offset+d.Size])
	}
	fnFiIdx := make([]int, len(img.Funcs))
	fnLinesIdx := make([]int, len(img.Funcs))
	fnDIEIdx := make([]int, len(img.Funcs))
	var dwarfRelocs []dwarfRelocSet
	for i, fn := range img.Funcs {
		data := marshalFuncInfo(fn)
		fnFiIdx[i] = len(defs)
		defs = append(defs, goSym{typ: kindSDATA, size: uint32(len(data))})
		defData = append(defData, data)

		name := fn.Name
		if !fn.Static {
			name = pkgPath + "." + name
		}

		// The DWARF symbols: the .debug_line state-machine program and the
		// subprogram DIE, both referencing the function by its non-package
		// index (package definitions, like cmd/asm's).
		lines, lrel := goobjDwarfLines(fn, fnNpIdx[i])
		fnLinesIdx[i] = len(defs)
		defs = append(defs, goSym{typ: kindSDWARFLINES, size: uint32(len(lines))})
		defData = append(defData, lines)
		die, drel := goobjDwarfInfo(fn, name, fnNpIdx[i])
		fnDIEIdx[i] = len(defs)
		defs = append(defs, goSym{typ: kindSDWARFFCN, size: uint32(len(die))})
		defData = append(defData, die)
		dwarfRelocs = append(dwarfRelocs,
			dwarfRelocSet{si: fnLinesIdx[i], relocs: lrel},
			dwarfRelocSet{si: fnDIEIdx[i], relocs: drel},
		)
	}

	// Resolve external symbol references (cross-package).  Build the
	// package index table and determine each external symbol's SymIdx
	// by reading the target package's export data.
	var extPkgTable []string
	var extPkgIdx map[string]int
	var extSymIdx map[string]int
	if len(img.Externals) > 0 {
		var err error
		extPkgTable, extPkgIdx, extSymIdx, err = resolveExternalSymbols(img.Externals)
		if err != nil {
			return nil, fmt.Errorf("GOOBJ emission: resolving external symbols: %w", err)
		}
	}

	// Relocations, per defined symbol in definition order (package defs,
	// then non-package defs).
	nsyms := len(defs) + len(nps)
	symRelocs := make([][]byte, nsyms) // flat 23-byte records
	for i, fn := range img.Funcs {
		si := len(defs) + fnNpIdx[i]
		for _, r := range fn.Relocs {
			typ, size := relocField(r)
			if r.External {
				// Split package-qualified name: "runtime·morestack" → runtime, morestack.
				pkg, name := splitQualified(r.Name)
				if pkg == "" {
					return nil, fmt.Errorf("GOOBJ emission: external symbol %q has no package prefix", r.Name)
				}
				pIdx, ok := extPkgIdx[pkg]
				if !ok {
					return nil, fmt.Errorf("GOOBJ emission: package %q not resolved", pkg)
				}
				sIdx, ok := extSymIdx[pkg+"·"+name]
				if !ok {
					return nil, fmt.Errorf("GOOBJ emission: symbol %s·%s not resolved", pkg, name)
				}
				var rec [23]byte
				binary.LittleEndian.PutUint32(rec[0:], uint32(int32(r.Off)))
				rec[4] = size // field width
				binary.LittleEndian.PutUint16(rec[5:], typ)
				binary.LittleEndian.PutUint64(rec[7:], uint64(r.Addend))
				binary.LittleEndian.PutUint32(rec[15:], uint32(pIdx))
				binary.LittleEndian.PutUint32(rec[19:], uint32(sIdx))
				symRelocs[si] = append(symRelocs[si], rec[:]...)
				continue
			}
			di, ok := defIdx[r.Name]
			if !ok {
				return nil, fmt.Errorf("GOOBJ emission: reference to unknown symbol %q", r.Name)
			}
			var rec [23]byte
			binary.LittleEndian.PutUint32(rec[0:], uint32(int32(r.Off)))
			rec[4] = size // field width
			binary.LittleEndian.PutUint16(rec[5:], typ)
			binary.LittleEndian.PutUint64(rec[7:], uint64(r.Addend))
			binary.LittleEndian.PutUint32(rec[15:], pkgIdxSelf)
			binary.LittleEndian.PutUint32(rec[19:], uint32(di))
			symRelocs[si] = append(symRelocs[si], rec[:]...)
		}
	}
	// The DWARF symbols' own relocations (the function address references).
	for _, ds := range dwarfRelocs {
		for _, r := range ds.relocs {
			var rec [23]byte
			binary.LittleEndian.PutUint32(rec[0:], uint32(r.off))
			rec[4] = r.siz
			binary.LittleEndian.PutUint16(rec[5:], r.typ)
			binary.LittleEndian.PutUint64(rec[7:], uint64(r.add))
			binary.LittleEndian.PutUint32(rec[15:], r.pkg)
			binary.LittleEndian.PutUint32(rec[19:], r.sym)
			symRelocs[ds.si] = append(symRelocs[ds.si], rec[:]...)
		}
	}

	// Aux entries per function: FuncInfo, the DWARF symbols, then the four
	// pc tables.  References into the non-package table use pkgIdxNone.
	symAux := make([][]byte, nsyms)
	for i := range img.Funcs {
		si := len(defs) + fnNpIdx[i]
		aux := func(typ uint8, pkg, idx uint32) {
			var rec [9]byte
			rec[0] = typ
			binary.LittleEndian.PutUint32(rec[1:], pkg)
			binary.LittleEndian.PutUint32(rec[5:], idx)
			symAux[si] = append(symAux[si], rec[:]...)
		}
		aux(auxFuncInfo, pkgIdxSelf, uint32(fnFiIdx[i]))
		aux(auxDwarfInfo, pkgIdxSelf, uint32(fnDIEIdx[i]))
		aux(auxDwarfLines, pkgIdxSelf, uint32(fnLinesIdx[i]))
		// The pc-table references are 0-based within the non-package
		// definitions; the loader adds the package-definition count itself.
		aux(auxPcsp, pkgIdxNone, uint32(pcIdx[i].sp))
		aux(auxPcfile, pkgIdxNone, uint32(pcIdx[i].file))
		aux(auxPcline, pkgIdxNone, uint32(pcIdx[i].line))
		aux(auxPcinline, pkgIdxNone, uint32(pcIdx[i].inl))
	}

	// The string table.  Absolute offsets: it starts right after the
	// 96-byte header (magic, fingerprint, flags, the 19 block offsets).
	const headerSize = 8 + 8 + 4 + 4*(blkEnd+1)
	strTab := []byte{}
	strOff := map[string]uint32{}
	addStr := func(s string) {
		if _, ok := strOff[s]; ok {
			return
		}
		strOff[s] = uint32(headerSize + len(strTab))
		strTab = append(strTab, s...)
	}
	addStr("")
	addStr(srcPath)
	for _, s := range defs {
		addStr(s.name)
	}
	for _, s := range nps {
		addStr(s.sym.name)
	}
	stringRef := func(b []byte, s string) []byte {
		b = binary.LittleEndian.AppendUint32(b, uint32(len(s)))
		return binary.LittleEndian.AppendUint32(b, strOff[s])
	}

	// Serialise the block bodies.
	var symdefBlk, npdefBlk []byte
	for _, s := range defs {
		symdefBlk = s.append(symdefBlk, strOff)
	}
	for _, s := range nps {
		npdefBlk = s.sym.append(npdefBlk, strOff)
	}

	// Package index table: index 0 is the dummy invalid package.
	// External packages follow, in pkgIdx order.
	for _, pkg := range extPkgTable {
		addStr(pkg)
	}
	pkgIdxBlk := stringRef(nil, "") // index 0: dummy
	for _, pkg := range extPkgTable {
		pkgIdxBlk = stringRef(pkgIdxBlk, pkg)
	}
	fileBlk := stringRef(nil, srcPath)

	var relocBlk, auxBlk, dataBlk []byte
	relocIdxBlk := make([]byte, 0, 4*(nsyms+1))
	auxIdxBlk := make([]byte, 0, 4*(nsyms+1))
	dataIdxBlk := make([]byte, 0, 4*(nsyms+1))
	var nr, na, nd uint32
	for si := 0; si < nsyms; si++ {
		relocIdxBlk = binary.LittleEndian.AppendUint32(relocIdxBlk, nr)
		auxIdxBlk = binary.LittleEndian.AppendUint32(auxIdxBlk, na)
		dataIdxBlk = binary.LittleEndian.AppendUint32(dataIdxBlk, nd)
		relocBlk = append(relocBlk, symRelocs[si]...)
		auxBlk = append(auxBlk, symAux[si]...)
		var d []byte
		if si < len(defData) {
			d = defData[si]
		} else {
			d = nps[si-len(defData)].data
		}
		dataBlk = append(dataBlk, d...)
		nr += uint32(len(symRelocs[si])) / 23
		na += uint32(len(symAux[si])) / 9
		nd += uint32(len(d))
	}
	relocIdxBlk = binary.LittleEndian.AppendUint32(relocIdxBlk, nr)
	auxIdxBlk = binary.LittleEndian.AppendUint32(auxIdxBlk, na)
	dataIdxBlk = binary.LittleEndian.AppendUint32(dataIdxBlk, nd)

	blocks := [blkEnd][]byte{
		blkPkgIdx:    pkgIdxBlk,
		blkFile:      fileBlk,
		blkSymdef:    symdefBlk,
		blkNonpkgdef: npdefBlk,
		blkRelocIdx:  relocIdxBlk,
		blkAuxIdx:    auxIdxBlk,
		blkDataIdx:   dataIdxBlk,
		blkReloc:     relocBlk,
		blkAux:       auxBlk,
		blkData:      dataBlk,
	}

	// Assemble the payload: header (offsets filled once known), string
	// table, blocks in order.
	payload := make([]byte, headerSize)
	copy(payload, goobjMagic)
	// The fingerprint stays zero, as cmd/asm leaves it.
	binary.LittleEndian.PutUint32(payload[16:], 4) // ObjFlagFromAssembly
	off := uint32(headerSize + len(strTab))
	for i := 0; i < blkEnd; i++ {
		binary.LittleEndian.PutUint32(payload[20+4*i:], off)
		off += uint32(len(blocks[i]))
	}
	binary.LittleEndian.PutUint32(payload[20+4*blkEnd:], off)
	payload = append(payload, strTab...)
	for _, blk := range blocks {
		payload = append(payload, blk...)
	}

	out := make([]byte, 0, len(pre)+len(payload))
	out = append(out, pre...)
	return append(out, payload...), nil
}

// marshalFuncInfo serialises a function's goobj.FuncInfo: sizes, flags,
// start line, the one-element file table and an empty inline tree.
func marshalFuncInfo(fn FuncLayout) []byte {
	flag := uint8(funcFlagAsm)
	if fn.SPWrite {
		flag |= funcFlagSPWrite
	}
	b := make([]byte, 0, 28)
	b = binary.LittleEndian.AppendUint32(b, uint32(fn.Args))
	b = binary.LittleEndian.AppendUint32(b, uint32(fn.Frame))
	b = append(b, 0, flag, 0, 0) // FuncID normal, flags, padding
	b = binary.LittleEndian.AppendUint32(b, uint32(int32(fn.Line)))
	b = binary.LittleEndian.AppendUint32(b, 1) // one file
	b = binary.LittleEndian.AppendUint32(b, 0) // file index 0
	b = binary.LittleEndian.AppendUint32(b, 0) // no inline tree
	return b
}

// pcValueFlat encodes a pc-value table holding v over the whole function.
// The pc deltas are in MinLC units (the runtime scales them by the
// architecture's minimum instruction length).
func pcValueFlat(v int32, size, minLC int) []byte {
	// The table is delta-encoded from an implicit value of -1: a varint
	// value delta, an unsigned pc delta to the end, and a zero terminator.
	out := binary.AppendVarint(nil, int64(v)+1)
	out = binary.AppendUvarint(out, uint64(size/minLC))
	return append(out, 0)
}

// pcspTable encodes the stack-adjustment table: the SP delta in effect at
// every pc, from the function's prologue and epilogue boundaries.
func pcspTable(fn FuncLayout, minLC int) []byte {
	if len(fn.Spadj) == 0 {
		return pcValueFlat(0, fn.Size, minLC)
	}
	pts := make([]SpadjStep, 0, len(fn.Spadj)+1)
	pts = append(pts, SpadjStep{PC: 0, Value: 0})
	pts = append(pts, fn.Spadj...)
	out := binary.AppendVarint(nil, int64(pts[0].Value)+1)
	cur, old := pts[0].PC, pts[0].Value
	for _, p := range pts[1:] {
		out = binary.AppendUvarint(out, uint64((p.PC-cur)/minLC))
		out = binary.AppendVarint(out, int64(p.Value-old))
		cur, old = p.PC, p.Value
	}
	out = binary.AppendUvarint(out, uint64((fn.Size-cur)/minLC))
	return append(out, 0)
}

// toolchainObjectPreamble returns the "go object ...\n!\n" header the
// installed go tool asm writes, captured by assembling a one-instruction
// probe.  The linker compares this string verbatim against its own, so it
// must come from the toolchain itself, not be reconstructed.
var (
	preambleOnce sync.Once
	preamble     []byte
	preambleErr  error
)

func toolchainObjectPreamble() ([]byte, error) {
	preambleOnce.Do(func() {
		goBin, err := exec.LookPath("go")
		if err != nil {
			preambleErr = fmt.Errorf("GOOBJ emission needs the Go toolchain: %w", err)
			return
		}
		dir, err := os.MkdirTemp("", "gasm-preamble")
		if err != nil {
			preambleErr = err
			return
		}
		defer os.RemoveAll(dir)
		src := filepath.Join(dir, "probe_amd64.s")
		if err := os.WriteFile(src, []byte("TEXT \u00b7x(SB), $0-0\n\tRET\n"), 0o644); err != nil {
			preambleErr = err
			return
		}
		obj := filepath.Join(dir, "probe.o")
		cmd := exec.Command(goBin, "tool", "asm", "-p", "probe", "-o", obj, src)
		cmd.Env = append(os.Environ(), "GOARCH=amd64")
		if out, err := cmd.CombinedOutput(); err != nil {
			preambleErr = fmt.Errorf("probing the assembler for the object header: %v\n%s", err, out)
			return
		}
		data, err := os.ReadFile(obj)
		if err != nil {
			preambleErr = err
			return
		}
		i := bytes.Index(data, []byte("\n!\n"))
		if i < 0 || !bytes.HasPrefix(data[i+3:], []byte(goobjMagic)) {
			preambleErr = fmt.Errorf("unrecognised assembler object layout")
			return
		}
		preamble = data[:i+3]
	})
	return preamble, preambleErr
}