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

// 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, which reserves index 0 for the dummy invalid
// package (cmd/internal/obj/sym.go: "0 is invalid index"; the loader's
// reader loop starts at 1), so package i sits at block index i+1 and its
// relocations carry i+1.  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 {
		// Block index 0 is the dummy invalid package; the first real
		// package starts at 1.
		pkgIdx[pkg.path] = i + 1
		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
// blocks.  The hashed blocks are kept raw: their symbols carry no names, only
// the loader needs their counts.
type goobjFile struct {
	strTab   []byte // string table, at headerSize + n
	symdef   []byte // blkSymdef raw block
	hashed64 []byte // blkHashed64def raw block
	hashed   []byte // blkHasheddef raw block
	npdef    []byte // blkNonpkgdef raw block
}

// loaderIndexBase returns the index the first nonpkgdef symbol occupies in the
// loader's per-object symbol array.  cmd/link lays the definition blocks out as
// symdef, hashed64def, hasheddef, nonpkgdef, nonpkgref (loader.go: preloadSyms
// fills r.syms in exactly that order, and resolve() indexes PkgIdxNone and
// cross-package SymIdx into it), so a symbol found in blkNonpkgdef carries the
// three leading blocks' symbol counts as its base.
func (f *goobjFile) loaderIndexBase() int {
	return len(f.symdef)/recSymSize + len(f.hashed64)/recSymSize + len(f.hashed)/recSymSize
}

// symbols returns the names of the symdef and nonpkgdef blocks in
// definition order.  Package definitions (blkSymdef) use fully-qualified
// names like "runtime.morestack"; non-package definitions (blkNonpkgdef)
// use bare names like "morestack".  For lookups by index prefer
// findSymbol: it adds the hashed blocks' count the loader's array
// interleaves between the two.
func (f *goobjFile) symbols() []string {
	return append(f.defNames(), f.npdefNames()...)
}

// findSymbol returns the index of a symbol within the loader's per-object
// symbol array, or -1 if not found.  It first tries the fully-qualified
// name (pkg.name), then the bare name (assembly objects store dotless
// names, e.g. runtime's "gogo", for symbols other packages reach through
// a linkname).
func (f *goobjFile) findSymbol(pkg, name string) int {
	base := f.loaderIndexBase()
	qualified := pkg + "." + name
	for i, s := range f.defNames() {
		if s == qualified {
			return i
		}
	}
	for i, s := range f.npdefNames() {
		if s == qualified {
			return base + i
		}
	}
	// Try bare name (for dotless assembly definitions).
	for i, s := range f.defNames() {
		if s == name {
			return i
		}
	}
	for i, s := range f.npdefNames() {
		if s == name {
			return base + 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)
}

// recSymSize is the size of one Sym record in the definition blocks
// (goobj.SymSize: stringRefSize + 2 + 1 + 1 + 1 + 4 + 4).
const recSymSize = 21

// 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 = recSymSize
	if len(block) < recSize {
		return nil
	}
	n := len(block) / recSize
	names := make([]string, 0, n)
	for i := range n {
		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 := range blkEnd + 1 {
		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),
		hashed64: blockSlice(payload, offs, blkHashed64def, blkHashed64def+1),
		hashed:   blockSlice(payload, offs, blkHasheddef, blkHasheddef+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.  The indices are 1-based
	// (0 is the dummy invalid package, written by the emitter itself), so
	// the table without the dummy is indexed one below.
	pkgTable = make([]string, len(pkgIdx1))
	for pkg, idx := range pkgIdx1 {
		pkgTable[idx-1] = 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),
// whose UTF-8 encoding is two bytes, so the search must be string-based:
// IndexByte would match only the second byte and leave the lead byte on
// the package path.  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 before, after, ok := strings.Cut(full, "\u00b7"); ok {
		return before, after
	}
	if before, after, ok := strings.Cut(full, "."); ok {
		return before, after
	}
	return "", full
}