// 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 [--corpus [dir]] [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.

With --corpus the audit changes shape: it assembles every .s file under the
given directory (default GOROOT/src) with the gasm encoder only, no
toolchain probing.  A file whose name carries a recognisable _arch suffix is
attempted for that architecture; a file without one is attempted for all
four, exactly as a GOARCH build would compile it.  The report gives the
per-architecture pass rates and the most common failure reasons, which drive
the encodability backlog by frequency rather than by table order.
`)
	corpus := fs.Bool("corpus", false, "assemble a corpus of .s files and report pass rates and failure reasons")
	if err := fs.Parse(args); err != nil {
		return err
	}
	if *corpus {
		return cmdAuditCorpus(fs.Args())
	}
	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)
}

// --- corpus audit -----------------------------------------------------------

// corpusTarget is one architecture row of the corpus report.
type corpusTarget struct {
	a    arch.Arch
	name string
}

// corpusTally accumulates one architecture's attempts over the corpus.
type corpusTally struct {
	attempted int
	assembled int
	reasons   map[string]int    // failure reason → count
	example   map[string]string // failure reason → one representative file
}

func (t *corpusTally) fail(path, reason string) {
	t.reasons[reason]++
	if t.example[reason] == "" {
		t.example[reason] = path
	}
}

// cmdAuditCorpus implements audit-instructions --corpus.
func cmdAuditCorpus(args []string) error {
	if len(args) > 1 {
		return fmt.Errorf("audit-instructions --corpus takes at most one directory argument")
	}
	root := ""
	if len(args) == 1 {
		root = args[0]
	} else {
		out, err := exec.Command("go", "env", "GOROOT").Output()
		if err != nil {
			return fmt.Errorf("locate GOROOT: %w", err)
		}
		root = filepath.Join(strings.TrimSpace(string(out)), "src")
	}
	stats, err := runCorpusAudit(root)
	if err != nil {
		return err
	}
	printCorpusStats(stats)
	return nil
}

// corpusStats is the outcome of one corpus audit run.
type corpusStats struct {
	root    string
	files   int
	generic int // files attempted for all four architectures
	full    int // files that assembled for every target architecture
	targets []corpusTarget
	tallies []*corpusTally
}

// runCorpusAudit assembles every .s file under root and returns the stats.
func runCorpusAudit(root string) (*corpusStats, error) {
	files, err := asmFiles(root)
	if err != nil {
		return nil, err
	}

	targets := []corpusTarget{
		{arch.AMD64, "amd64"},
		{arch.ARM64, "arm64"},
		{arch.RISCV, "riscv64"},
		{arch.LOONG64, "loong64"},
	}
	tallies := make([]*corpusTally, len(targets))
	for i := range tallies {
		tallies[i] = &corpusTally{reasons: map[string]int{}, example: map[string]string{}}
	}
	// full is the north-star number: a file counts when every architecture
	// its name allows assembles it.
	full, generic := 0, 0

	for _, path := range files {
		src, err := readSource(path)
		if err != nil {
			return nil, err
		}
		f, errs := parser.Parse(path, src)

		var wanted []int // indexes into targets
		if a := arch.FromFilename(path); a != arch.Unknown {
			for i, tg := range targets {
				if tg.a == a {
					wanted = append(wanted, i)
				}
			}
		} else {
			generic++
			for i := range targets {
				wanted = append(wanted, i)
			}
		}

		ok := true
		for _, i := range wanted {
			tg, t := targets[i], tallies[i]
			t.attempted++
			var err error
			if len(errs) > 0 {
				err = errs[0] // a parse failure is a failure for every target
			} else {
				_, err = assembleFile(tg.a, f)
			}
			if err != nil {
				ok = false
				t.fail(path, corpusReason(err))
				continue
			}
			t.assembled++
		}
		if ok && len(wanted) > 0 {
			full++
		}
	}

	return &corpusStats{
		root:    root,
		files:   len(files),
		generic: generic,
		full:    full,
		targets: targets,
		tallies: tallies,
	}, nil
}

// printCorpusStats renders the corpus audit report.
func printCorpusStats(s *corpusStats) {
	fmt.Printf("corpus %s: %d files (%d generic, attempted for all architectures)\n", s.root, s.files, s.generic)
	fmt.Printf("  assemble for every target architecture: %d (%.1f%%)\n", s.full, 100*float64(s.full)/float64(max(s.files, 1)))
	for i, tg := range s.targets {
		t := s.tallies[i]
		fmt.Printf("  %s: %d/%d attempted\n", tg.name, t.assembled, t.attempted)
		for _, r := range topReasons(t) {
			fmt.Printf("    %4d  %s\n", t.reasons[r], r)
			fmt.Printf("         e.g. %s\n", t.example[r])
		}
	}
}

// corpusReason buckets an assembly or parse failure for the histogram.
func corpusReason(err error) string {
	msg := err.Error()
	switch {
	case strings.Contains(msg, "unsupported"), strings.Contains(msg, "cannot encode"):
		return "instruction not encodable"
	case strings.Contains(msg, "undefined label"):
		return "undefined label"
	case strings.Contains(msg, "undefined symbol"), strings.Contains(msg, "external symbol"), strings.Contains(msg, "file-level assembly"):
		return "undefined symbol or external"
	case strings.Contains(msg, "operand"), strings.Contains(msg, "operand form"):
		return "unsupported operand form"
	default:
		return "other: " + firstLine(msg)
	}
}

// topReasons returns at most five reasons, most frequent first.
func topReasons(t *corpusTally) []string {
	type kv struct {
		k string
		n int
	}
	var kvs []kv
	for k, n := range t.reasons {
		kvs = append(kvs, kv{k, n})
	}
	slices.SortFunc(kvs, func(a, b kv) int { return b.n - a.n })
	if len(kvs) > 5 {
		kvs = kvs[:5]
	}
	out := make([]string, len(kvs))
	for i, kv := range kvs {
		out[i] = kv.k
	}
	return out
}

// firstLine returns the first line of an error message, truncated.
func firstLine(msg string) string {
	if i := strings.IndexByte(msg, '\n'); i >= 0 {
		msg = msg[:i]
	}
	if len(msg) > 80 {
		msg = msg[:80]
	}
	return msg
}