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

package verify

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

// GroundTruth assembles the given .s file with the Go toolchain's own
// assembler and returns the machine code bytes for each TEXT function,
// keyed by the function's short name (the part after the middle dot).
// This is the universal oracle: any file that `go tool asm` accepts can
// be verified, with no hand-written reference.
//
// For RISC-V sources the assembler is invoked with GOARCH=riscv64;
// the caller must set the architecture via GroundTruthArch.
func GroundTruth(path string) (map[string][]byte, error) {
	return groundTruthArch(path, "")
}

// GroundTruthRISCV assembles the given .s file with the Go toolchain in
// RISC-V cross-assembly mode (GOARCH=riscv64).
func GroundTruthRISCV(path string) (map[string][]byte, error) {
	return groundTruthArch(path, "riscv64")
}

func groundTruthArch(path, goarch string) (map[string][]byte, error) {
	goroot := runtime.GOROOT()
	asmBin := filepath.Join(goroot, "pkg", "tool", runtime.GOOS+"_"+runtime.GOARCH, "asm")
	if _, err := os.Stat(asmBin); err != nil {
		return nil, fmt.Errorf("verify: go tool asm not found at %s: %w", asmBin, err)
	}
	includeDir := filepath.Join(goroot, "pkg", "include")

	tmpDir, err := os.MkdirTemp("", "gasm-verify-*")
	if err != nil {
		return nil, fmt.Errorf("verify: tempdir: %w", err)
	}
	defer os.RemoveAll(tmpDir)
	objPath := filepath.Join(tmpDir, "out.o")

	base := filepath.Base(path)
	pkg := strings.TrimSuffix(base, ".s")
	pkg = strings.TrimSuffix(pkg, "_amd64")
	pkg = strings.TrimSuffix(pkg, "_riscv64")

	cmd := exec.Command(asmBin, "-I", includeDir, "-p", pkg, "-o", objPath, path)
	if goarch != "" {
		cmd.Env = append(os.Environ(), "GOARCH="+goarch)
	}
	if out, err := cmd.CombinedOutput(); err != nil {
		return nil, fmt.Errorf("verify: go tool asm (%s): %w\n%s", goarch, err, out)
	}

	objData, err := os.ReadFile(objPath)
	if err != nil {
		return nil, fmt.Errorf("verify: read object: %w", err)
	}
	return extractGOOBJCode(objData)
}

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

const goobjMagic = "\x00go120ld"

// extractGOOBJCode parses a GOOBJ payload and returns the code bytes for
// each non-package STEXT symbol (the functions).
func extractGOOBJCode(data []byte) (map[string][]byte, error) {
	// Find the GOOBJ header (after the "go object ..." preamble).
	i := bytes.Index(data, []byte(goobjMagic))
	if i < 0 {
		return nil, fmt.Errorf("verify: no GOOBJ magic in object file")
	}
	b := data[i:]
	le := binary.LittleEndian

	// Read block offsets (20 bytes into the header: 4 magic + 8 go version
	// + 8 experiment = 20, then blkEnd+1 uint32 offsets).
	var offs [blkEnd + 1]uint32
	for j := 0; j <= blkEnd; j++ {
		offs[j] = le.Uint32(b[20+4*j:])
	}
	blk := func(idx int) []byte { return b[offs[idx]:offs[idx+1]] }

	// Parse non-package symbol definitions (blkNonpkgdef): each entry is
	// 21 bytes: [nameLen:4][nameOff:4][abi:2][type:1][flag:1][flag2:1][size:4][align:4].
	const symSize = 21
	nonpkg := blk(blkNonpkgdef)
	nSyms := len(nonpkg) / symSize

	// Data index (blkDataIdx): one uint32 per defined symbol across ALL
	// definition blocks (blkSymdef + blkHashed64def + blkHasheddef +
	// blkNonpkgdef), in that order.  We need the offset for the nonpkg
	// symbols, which come last.
	dataIdx := blk(blkDataIdx)
	dataBlk := blk(blkData)

	// Count symbols in the preceding definition blocks.
	preceding := 0
	for _, bi := range []int{blkSymdef, blkHashed64def, blkHasheddef} {
		preceding += len(blk(bi)) / symSize
	}

	// Symbol name offsets in the GOOBJ symbol table are absolute byte
	// offsets from the start of the GOOBJ payload (the magic).
	readStr := func(off, ln uint32) string {
		if int(off+ln) > len(b) {
			return ""
		}
		return string(b[off : off+ln])
	}

	result := make(map[string][]byte)
	const kindSTEXT = 1
	for s := 0; s < nSyms; s++ {
		x := nonpkg[s*symSize:]
		nameLen := le.Uint32(x[0:])
		nameOff := le.Uint32(x[4:])
		typ := x[10]
		size := le.Uint32(x[13:])

		if typ != kindSTEXT || size == 0 {
			continue
		}
		name := readStr(nameOff, nameLen)
		// Strip the package prefix (everything up to and including the
		// last middle dot or period-dot).
		name = stripPkg(name)

		// Data offset from the index (nonpkg symbols follow the preceding blocks).
		diIdx := preceding + s
		if (diIdx+1)*4 > len(dataIdx) {
			continue
		}
		dOff := le.Uint32(dataIdx[diIdx*4:])
		if int(dOff+size) > len(dataBlk) {
			continue
		}
		code := make([]byte, size)
		copy(code, dataBlk[dOff:dOff+size])
		result[name] = code
	}
	return result, nil
}

// stripPkg removes the package path prefix from a symbol name, leaving
// just the function name.  "pkg/path·FuncName" → "FuncName".
func stripPkg(name string) string {
	if i := strings.LastIndex(name, "\u00B7"); i >= 0 {
		return name[i+len("\u00B7"):]
	}
	if i := strings.LastIndex(name, "\"."); i >= 0 {
		return name[i+2:]
	}
	if i := strings.LastIndex(name, "."); i >= 0 {
		return name[i+1:]
	}
	return name
}