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

package main

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"slices"
	"strconv"
	"strings"

	"sourcedock.dev/petrbalvin/gasm-devkit/arch"
	"sourcedock.dev/petrbalvin/gasm-devkit/asm"
	"sourcedock.dev/petrbalvin/gasm-devkit/parser"
)

// cmdAuditInstructions cross-checks a gasm encoder against the Go toolchain's
// own assembler, probed black-box: every mnemonic in the gasm table is offered
// to go tool asm in its bare form, and a mnemonic counts as known to Go when
// the error is anything but "unrecognized instruction" (a wrong-shape error
// still proves the mnemonic exists in Go's tables).  The audit answers three
// questions at a glance:
//
//   - which mnemonics gasm can encode that go tool asm does not know
//     (superset encodings, usable only through the gasm goobj path);
//   - which mnemonics the architecture table knows but the encoder cannot
//     emit yet (the implementation backlog);
//   - which mnemonics go tool asm knows that gasm cannot encode (feature
//     gaps).
//
// The amd64 derived families (Jcc, CMOVcc, SETcc) exist on both sides by
// construction and are excluded from the diff; the other architectures list
// their conditional branches outright.
func cmdAuditInstructions(args []string) error {
	fs := newCommand("audit-instructions", "gasm audit-instructions [amd64|arm64|riscv64|loong64]", `
Compare the gasm encoder for the given architecture (default amd64) against
go tool asm and print the diff: superset encodings (gasm-only, shippable via
gasm asm --format goobj), known-but-unencodable names (the backlog) and go-
only names (feature gaps).  The Go side is probed black-box with a battery
of bare mnemonics, so the audit tracks whatever toolchain `+"`go env GOROOT`"+`
provides.
`)
	if err := fs.Parse(args); err != nil {
		return err
	}
	archName := "amd64"
	switch n := len(fs.Args()); {
	case n > 1:
		return fmt.Errorf("audit-instructions takes at most one architecture argument")
	case n == 1:
		archName = strings.ToLower(fs.Arg(0))
	}
	a, err := auditArch(archName)
	if err != nil {
		return err
	}

	tab := arch.ForArch(a)
	var names []string
	seen := map[string]bool{}
	for _, in := range tab.Instructions() {
		name := strings.ToUpper(in.Name)
		if a == arch.AMD64 && derivedFamily(name) || seen[name] {
			continue
		}
		seen[name] = true
		names = append(names, name)
	}

	goKnown, err := probeGoAsm(goarchName(a), names)
	if err != nil {
		return err
	}

	var superset, backlog, shared []string
	for _, name := range names {
		switch {
		case !gasmEncodable(a, name):
			backlog = append(backlog, name)
		case !goKnown[name]:
			superset = append(superset, name)
		default:
			shared = append(shared, name)
		}
	}
	// GO-ONLY is not enumerable by probing: Go's table is only visible
	// through names we already know, so nothing can be reported there.

	slices.Sort(superset)
	slices.Sort(backlog)
	slices.Sort(shared)

	w := os.Stdout
	fmt.Fprintf(w, "gasm table (%s, families excluded): %d mnemonics\n", archName, len(names))
	fmt.Fprintf(w, "gasm encodable: %d   go tool asm recognized: %d\n", len(shared)+len(superset), countTrue(goKnown))
	fmt.Fprintf(w, "shared: %d\n", len(shared))
	fmt.Fprintf(w, "\nSuperset encodings (gasm-only; ship via gasm asm --format goobj):\n")
	for _, n := range superset {
		fmt.Fprintf(w, "  %s\n", n)
	}
	fmt.Fprintf(w, "\nKnown but not encodable (backlog):\n")
	for _, n := range backlog {
		fmt.Fprintf(w, "  %s\n", n)
	}
	fmt.Fprintf(w, "\nGo-only names cannot be enumerated by probing; extend the gasm\n")
	fmt.Fprintf(w, "table from the Go release notes when a new instruction family ships.\n")
	return nil
}

// auditArch resolves the audit's architecture argument.
func auditArch(name string) (arch.Arch, error) {
	switch strings.ToLower(name) {
	case "amd64":
		return arch.AMD64, nil
	case "arm64":
		return arch.ARM64, nil
	case "riscv64", "riscv":
		return arch.RISCV, nil
	case "loong64", "loong":
		return arch.LOONG64, nil
	}
	return arch.Unknown, fmt.Errorf("unknown architecture %q: want amd64, arm64, riscv64 or loong64", name)
}

// goarchName maps an arch identifier onto its GOARCH spelling.
func goarchName(a arch.Arch) string {
	switch a {
	case arch.ARM64:
		return "arm64"
	case arch.RISCV:
		return "riscv64"
	case arch.LOONG64:
		return "loong64"
	}
	return "amd64"
}

func countTrue(m map[string]bool) int {
	n := 0
	for _, v := range m {
		if v {
			n++
		}
	}
	return n
}

// derivedFamily reports whether a mnemonic belongs to a family both
// assemblers construct from condition codes rather than list exhaustively
// (JEQ/CMOVLGT/SETNE and friends).  Such names never probe cleanly, so
// including them in the diff would be noise.  amd64 only: the other
// architectures list their conditional branches outright.
func derivedFamily(name string) bool {
	if strings.HasPrefix(name, "J") && name != "JMP" && name != "JMPQ" {
		return true
	}
	if strings.HasPrefix(name, "CMOV") || strings.HasPrefix(name, "SET") {
		return true
	}
	return false
}

var unrecognizedRe = regexp.MustCompile(`unrecognized instruction`)

// probeGoAsm feeds every mnemonic to go tool asm in one generated file and
// classifies the diagnostics.  "Unrecognized instruction" is a parse-stage
// verdict on the mnemonic alone, so a single bare-instruction probe per
// mnemonic decides recognition; the combined file still reports every line's
// error even when others fail.
func probeGoAsm(goarch string, names []string) (map[string]bool, error) {
	dir, err := os.MkdirTemp("", "gasm-audit")
	if err != nil {
		return nil, err
	}
	defer os.RemoveAll(dir)

	var sb strings.Builder
	sb.WriteString("TEXT ·probe(SB), 4, $0\n\tRET\n")
	lineMnemonic := map[int]string{}
	line := 3
	for _, name := range names {
		fmt.Fprintf(&sb, "TEXT ·p%s%d(SB), 4, $0\n", sanitize(name), line)
		sb.WriteString("\t" + name + "\n\tRET\n")
		lineMnemonic[line+1] = name // the instruction line, after TEXT
		line += 3
	}
	probePath := filepath.Join(dir, "probe.s")
	if err := os.WriteFile(probePath, []byte(sb.String()), 0o644); err != nil {
		return nil, err
	}

	toolDir, err := exec.Command("go", "env", "GOTOOLDIR").Output()
	if err != nil {
		return nil, fmt.Errorf("go env GOTOOLDIR: %w", err)
	}
	asmBin := filepath.Join(strings.TrimSpace(string(toolDir)), "asm")
	if _, err := os.Stat(asmBin); err != nil {
		return nil, fmt.Errorf("go tool asm not found at %s", asmBin)
	}
	cmd := exec.Command(asmBin, "-p", "probe", "-o", filepath.Join(dir, "probe.o"), probePath)
	cmd.Env = append(os.Environ(), "GOARCH="+goarch, "GOOS="+runtime.GOOS)
	out, _ := cmd.CombinedOutput()

	result := map[string]bool{}
	for _, name := range names {
		result[name] = true // no news = the name parsed fine
	}
	reParse := regexp.MustCompile(`probe\.s:(\d+):`)
	for l := range strings.SplitSeq(string(out), "\n") {
		m := reParse.FindStringSubmatch(l)
		if m == nil {
			continue
		}
		lineNo, err := strconv.Atoi(m[1])
		if err != nil {
			continue
		}
		if name, ok := lineMnemonic[lineNo]; ok && unrecognizedRe.MatchString(l) {
			result[name] = false
		}
	}
	return result, nil
}

// probeShapes lists representative operand shapes for the encodability
// probe.  The assemblers report an unknown mnemonic and a known mnemonic
// with no supported form alike ("unsupported <arch> instruction"), so only
// a shape that assembles cleanly counts, and the backlog over-approximates:
// a name whose real forms the battery misses lands there.  amd64 keeps its
// exact table-driven check.
func probeShapes(a arch.Arch) []string {
	switch a {
	case arch.ARM64:
		return []string{
			"X0, X1, X2", "X0, X1", "X0", "$1, X0", "X0, (X1)", "(X0), X1",
			"X0, (X1, 8)", "(SP), X0", "F0, F1, F2", "F0, F1", "F0",
			"V0.B16, V1.B16, V2.B16", "p2", "X0, p2", "X0, X1, p2",
			// The conditional select family spells the condition first
			// and takes R register spellings.
			"EQ, R0, R1, R2", "EQ, R0, R1", "EQ, R0",
			"GE, F0, F1, F2", "NE, F0, F1, $0",
		}
	case arch.RISCV:
		return []string{
			"X5, X6, X7", "X5, X6", "X5", "$1, X5", "X5, (X6)", "$1, X5, X6",
			"(X5), X6", "F0, F1, F2", "F0, F1", "p2", "X1, p2", "X0, p2",
			"X5, X6, p2", "p2(SB)",
		}
	case arch.LOONG64:
		return []string{
			"R4, R5, R6", "R4, R5", "R4", "$1, R4", "R4, (R5)", "(R4), R5",
			"F0, F1, F2", "F0, F1", "p2", "R1, p2", "R4, p2",
			"$1, R4, R5, R6", "$65536, R4", "R4, R5, p2", "p2(SB)",
		}
	}
	return nil
}

// gasmEncodable reports whether the gasm encoder for a can emit the
// mnemonic, decided by trial assembly over the shape battery.
func gasmEncodable(a arch.Arch, name string) bool {
	switch a {
	case arch.ARM64, arch.RISCV, arch.LOONG64:
	default:
		return asm.Encodable(name)
	}
	for _, shape := range probeShapes(a) {
		if gasmAssembles(a, name, shape) {
			return true
		}
	}
	return false
}

// gasmAssembles reports whether a one-instruction probe file containing name
// with the given operand shape assembles without error.
func gasmAssembles(a arch.Arch, name, shape string) bool {
	src := "TEXT ·p(SB), NOSPLIT, $0\n\t" + name
	if shape != "" {
		src += " " + shape
	}
	src += "\n\tRET\np2:\n\tRET\n"
	f, errs := parser.Parse("probe.s", src)
	if len(errs) > 0 {
		return false
	}
	var err error
	switch a {
	case arch.ARM64:
		_, err = asm.AssembleFileARM64(f)
	case arch.RISCV:
		_, err = asm.AssembleFileRISCV(f)
	case arch.LOONG64:
		_, err = asm.AssembleFileLOONG64(f)
	}
	return err == nil
}

// sanitize makes a mnemonic safe for use in a Go symbol name.
func sanitize(name string) string {
	return strings.NewReplacer(".", "_", "$", "_").Replace(name)
}