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

// readGOOBJSymbols reads the GOOBJ symbol definitions from a compiled Go
// package's export file.  The file is an ar archive containing a __.PKGDEF
// member whose payload is the "go object ...\n!\n" preamble followed by the
// GOOBJ data.  The function returns the symbol names in definition order
// (the order they appear in blkSymdef), which matches the SymIdx the linker
// expects for cross-package references.
func readGOOBJSymbols(exportPath string) ([]string, error) {
	data, err := os.ReadFile(exportPath)
	if err != nil {
		return nil, err
	}
	goobj, err := extractGOOBJ(data)
	if err != nil {
		return nil, fmt.Errorf("%s: %w", exportPath, err)
	}
	return goobj.symbols(), nil
}

// exportPath returns the export file path for a given import path by running
// "go list -export".  The result is cached so repeated calls for the same
// package are fast.
func exportPath(importPath string) (string, error) {
	cmd := exec.Command("go", "list", "-json", "-export", importPath)
	out, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("go list %s: %w", importPath, err)
	}
	// Quick JSON extraction: find "Export": "…"
	const key = `"Export": "`
	i := bytes.Index(out, []byte(key))
	if i < 0 {
		return "", fmt.Errorf("go list %s: no Export field", importPath)
	}
	start := i + len(key)
	end := bytes.IndexByte(out[start:], '"')
	if end < 0 {
		return "", fmt.Errorf("go list %s: malformed Export field", importPath)
	}
	return string(out[start : start+end]), nil
}

// resolveExternalGOOBJ resolves a set of external symbol references into
// (package index, symbol index) pairs suitable for GOOBJ emission.
//
// refs maps package import paths to the symbol names referenced from that
// package.  The returned pkgIdx maps each import path to its position in
// the blkPkgIdx table (0-based), and symIdx gives each symbol's index within
// its package.
func resolveExternalGOOBJ(refs map[string][]string) (pkgIdx map[string]int, symIdx map[string]int, err error) {
	pkgIdx = make(map[string]int, len(refs))
	symIdx = make(map[string]int)

	// Assign package indices in sorted order for determinism.
	packages := sortedPkgRefs(refs)

	for i, pkg := range packages {
		pkgIdx[pkg.path] = i
		exp, err := exportPath(pkg.path)
		if err != nil {
			return nil, nil, err
		}
		data, err := os.ReadFile(exp)
		if err != nil {
			return nil, nil, err
		}
		gobj, err := extractGOOBJ(data)
		if err != nil {
			return nil, nil, fmt.Errorf("%s: %w", pkg.path, err)
		}
		for _, name := range pkg.syms {
			idx := gobj.findSymbol(pkg.path, name)
			if idx < 0 {
				return nil, nil, fmt.Errorf("symbol %s·%s not found in export data of %s", pkg.path, name, pkg.path)
			}
			symIdx[pkg.path+"·"+name] = idx
		}
	}
	return pkgIdx, symIdx, nil
}

type pkgRef struct {
	path string
	syms []string
}

func sortedPkgRefs(refs map[string][]string) []pkgRef {
	var pkgs []pkgRef
	for pkg, syms := range refs {
		pkgs = append(pkgs, pkgRef{pkg, syms})
	}
	// Simple insertion sort — the list is tiny (usually 1–3 packages).
	for i := 1; i < len(pkgs); i++ {
		for j := i; j > 0 && pkgs[j-1].path > pkgs[j].path; j-- {
			pkgs[j-1], pkgs[j] = pkgs[j], pkgs[j-1]
		}
	}
	return pkgs
}

// extractGOOBJ finds the GOOBJ data in an ar archive and returns a parsed
// goobjFile.  The archive member _go_.o contains the "go object …\n!\n"
// preamble followed by the GOOBJ payload; __.PKGDEF is the compiler export
// data (type information) and is not the GOOBJ object.
func extractGOOBJ(data []byte) (*goobjFile, error) {
	if len(data) < 8 || string(data[:8]) != "!<arch>\n" {
		return nil, fmt.Errorf("not an ar archive")
	}
	pos := 8
	for pos+60 <= len(data) {
		hdr := data[pos : pos+60]
		pos += 60

		// Parse ar header fields.
		name := strings.TrimRight(string(hdr[:16]), " /")
		size := parseArDecimal(hdr[48:58])
		if size < 0 {
			return nil, fmt.Errorf("invalid ar header: bad size")
		}
		if pos+size > len(data) {
			return nil, fmt.Errorf("ar entry %q extends past end of file", name)
		}
		body := data[pos : pos+size]
		pos += size
		// ar pads to even bytes.
		if pos%2 != 0 {
			pos++
		}

		if name == "_go_.o" {
			return parseGOOBJ(body)
		}
	}
	return nil, fmt.Errorf("archive contains no _go_.o member")
}

// parseArDecimal parses a decimal number from a space-padded field.
func parseArDecimal(b []byte) int {
	v := 0
	for _, c := range b {
		if c == ' ' {
			continue
		}
		if c < '0' || c > '9' {
			return -1
		}
		v = v*10 + int(c-'0')
	}
	return v
}

// goobjFile is a parsed GOOBJ file: the string table and the symbol-definition
// block.
type goobjFile struct {
	strTab []byte // string table, at headerSize + n
	symdef []byte // blkSymdef raw block
	npdef  []byte // blkNonpkgdef raw block
}

// symbols returns all symbol names in definition order by scanning the
// symdef and nonpkgdef blocks and resolving each name through the string
// table.  Package definitions (blkSymdef) use fully-qualified names like
// "runtime.morestack"; non-package definitions (blkNonpkgdef) use bare
// names like "morestack".  This combined list matches the index the
// linker expects for cross-package references.
func (f *goobjFile) symbols() []string {
	return append(f.defNames(), f.npdefNames()...)
}

// findSymbol returns the index of a symbol within the combined symbol list,
// or -1 if not found.  It first tries the fully-qualified name (pkg.name),
// then the bare name.
func (f *goobjFile) findSymbol(pkg, name string) int {
	qualified := pkg + "." + name
	syms := f.symbols()
	for i, s := range syms {
		if s == qualified {
			return i
		}
	}
	// Try bare name (for non-package definitions).
	for i, s := range syms {
		if s == name {
			return i
		}
	}
	return -1
}

// defNames returns names from blkSymdef only.
func (f *goobjFile) defNames() []string {
	return f.readSymNames(f.symdef)
}

// npdefNames returns names from blkNonpkgdef.
func (f *goobjFile) npdefNames() []string {
	return f.readSymNames(f.npdef)
}

// readSymNames reads symbol names from a symdef/nonpkgdef block.  Each record
// is 21 bytes: nameLen (u32), nameOff (u32), abi (u16), typ, flag, flag2,
// size (u32), align (u32).  nameOff is an absolute offset into the string
// table.
func (f *goobjFile) readSymNames(block []byte) []string {
	const recSize = 21
	if len(block) < recSize {
		return nil
	}
	n := len(block) / recSize
	names := make([]string, 0, n)
	for i := 0; i < n; i++ {
		rec := block[i*recSize : (i+1)*recSize]
		nameLen := binary.LittleEndian.Uint32(rec[0:4])
		nameOff := binary.LittleEndian.Uint32(rec[4:8])
		// nameOff is an absolute offset into the GOOBJ payload.  The string
		// table we have starts at goobjHeaderSize, so we subtract that.
		if nameOff < goobjHeaderSize {
			continue
		}
		relOff := nameOff - goobjHeaderSize
		if relOff >= uint32(len(f.strTab)) || relOff+nameLen > uint32(len(f.strTab)) {
			continue
		}
		names = append(names, string(f.strTab[relOff:relOff+nameLen]))
	}
	return names
}

const goobjHeaderSize = 8 + 8 + 4 + 4*(blkEnd+1) // magic + fingerprint + flags + 19 block offsets

// parseGOOBJ parses a raw GOOBJ payload (the data after the "\n!\n" preamble).
func parseGOOBJ(data []byte) (*goobjFile, error) {
	// Find the "\n!\n" separator.
	sep := []byte("\n!\n")
	i := bytes.Index(data, sep)
	if i < 0 {
		// Maybe the data has no preamble (e.g. a raw .o file).
		i = -3 // treat as if preamble starts before the data
	}
	payload := data[i+len(sep):]

	if len(payload) < goobjHeaderSize {
		return nil, fmt.Errorf("GOOBJ payload too short (%d bytes)", len(payload))
	}
	if string(payload[:8]) != goobjMagic {
		return nil, fmt.Errorf("bad GOOBJ magic: %q", payload[:8])
	}

	// Read block offsets.  The header layout is:
	//   [0:8]   magic
	//   [8:16]  fingerprint
	//   [16:20] flags
	//   [20:96] 19 × uint32 offsets
	var offs [blkEnd + 1]uint32
	for i := 0; i <= blkEnd; i++ {
		offs[i] = binary.LittleEndian.Uint32(payload[20+4*i:])
	}
	// The string table lives at headerSize.
	strTabStart := uint32(goobjHeaderSize)

	f := &goobjFile{
		strTab: payload[strTabStart:offs[0]],
		symdef: blockSlice(payload, offs, blkSymdef, blkSymdef+1),
		npdef:  blockSlice(payload, offs, blkNonpkgdef, blkNonpkgdef+1),
	}
	return f, nil
}

// blockSlice extracts a block from the payload using its offset pair.
func blockSlice(payload []byte, offs [blkEnd + 1]uint32, start, end int) []byte {
	if start < 0 || end > blkEnd || offs[end] < offs[start] {
		return nil
	}
	beg := offs[start]
	fin := offs[end]
	if int(fin) > len(payload) || int(beg) > int(fin) {
		return nil
	}
	return payload[beg:fin]
}

// resolveExternalSymbols is the high-level entry point for GOOBJ emission.
// Given a list of external symbol names (e.g. ["runtime·morestack",
// "runtime·g0"]), it returns the package-index table entries and a map from
// full symbol name to GOOBJ {pkgIdx, symIdx}.
//
// The package table entries should be written into blkPkgIdx, and the
// returned indices should replace pkgIdxSelf / placeholder values in the
// relocation records.
func resolveExternalSymbols(externals []string) (pkgTable []string, pkgIdxMap map[string]int, symIdxMap map[string]int, err error) {
	// Group references by package.
	refs := make(map[string]map[string]bool)
	for _, full := range externals {
		pkg, name := splitQualified(full)
		if refs[pkg] == nil {
			refs[pkg] = make(map[string]bool)
		}
		refs[pkg][name] = true
	}

	// Convert maps to slices.
	r := make(map[string][]string, len(refs))
	for pkg, names := range refs {
		for name := range names {
			r[pkg] = append(r[pkg], name)
		}
	}

	pkgIdx1, symIdx1, err := resolveExternalGOOBJ(r)
	if err != nil {
		return nil, nil, nil, err
	}

	// Build the package table in pkgIdx order.
	pkgTable = make([]string, len(pkgIdx1))
	for pkg, idx := range pkgIdx1 {
		pkgTable[idx] = pkg
	}

	return pkgTable, pkgIdx1, symIdx1, nil
}

// splitQualified splits a qualified Go symbol name (pkgpath·name) into its
// package path and local name.  The separator is the middle dot (U+00B7).
// If no separator is found, the symbol is assumed to be in the current
// package (empty pkg).
func splitQualified(full string) (pkg, name string) {
	if idx := strings.IndexByte(full, '\u00b7'); idx >= 0 {
		return full[:idx], full[idx+len("\u00b7"):]
	}
	if idx := strings.IndexByte(full, '.'); idx >= 0 {
		return full[:idx], full[idx+1:]
	}
	return "", full
}