// 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, and the pc-value tables (pcsp, pcfile,
// pcline, pcinline).  DWARF and the implicit funcdata symbols are omitted;
// the linker fills their defaults.

// 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
)

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

// 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)
}

// 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) {
	if pkgPath == "" {
		return nil, fmt.Errorf("GOOBJ emission requires a package path (-p)")
	}
	pre, err := toolchainObjectPreamble()
	if err != nil {
		return nil, err
	}

	// The symbol tables.  Package definitions: the GLOBL symbols, then one
	// anonymous FuncInfo symbol per function.  Non-package definitions: the
	// pc-value tables and the functions themselves, as cmd/asm lays them
	// out.  defIdx maps a GLOBL's bare name to its definition index for the
	// relocations; fnNpIdx maps a function to its non-package index.
	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))
	for i := range img.Funcs {
		data := marshalFuncInfo(img.Funcs[i])
		fnFiIdx[i] = len(defs)
		defs = append(defs, goSym{typ: kindSDATA, size: uint32(len(data))})
		defData = append(defData, data)
	}

	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), &pcIdx[i].sp},
			{pcValueFlat(0, fn.Size), &pcIdx[i].file},
			{pcValueFlat(int32(fn.Line), fn.Size), &pcIdx[i].line},
			{pcValueFlat(-1, fn.Size), &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 {
			// The linker writes the resolved displacement into the field;
			// leave it zero, as cmd/asm's object does.
			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,
		})
	}

	// Relocations, per defined symbol in definition order (package defs,
	// then non-package defs).  Only file-local GLOBL references resolve;
	// external symbols need the import machinery of a later increment.
	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 {
			if r.External {
				return nil, fmt.Errorf("GOOBJ emission: external symbol %q is not supported yet", r.Name)
			}
			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] = 4 // field width
			binary.LittleEndian.PutUint16(rec[5:], relocPCRel)
			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[:]...)
		}
	}

	// Aux entries per function: FuncInfo, 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(auxPcsp, pkgIdxNone, uint32(len(defs)+pcIdx[i].sp))
		aux(auxPcfile, pkgIdxNone, uint32(len(defs)+pcIdx[i].file))
		aux(auxPcline, pkgIdxNone, uint32(len(defs)+pcIdx[i].line))
		aux(auxPcinline, pkgIdxNone, uint32(len(defs)+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)
	}
	pkgIdxBlk := stringRef(nil, "") // index 0: the dummy invalid package
	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.
func pcValueFlat(v int32, size 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))
	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) []byte {
	if len(fn.Spadj) == 0 {
		return pcValueFlat(0, fn.Size)
	}
	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))
		out = binary.AppendVarint(out, int64(p.Value-old))
		cur, old = p.PC, p.Value
	}
	out = binary.AppendUvarint(out, uint64(fn.Size-cur))
	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
}