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

package main

import (
	"bytes"
	"io"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"testing"

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

const clean = "#include \"textflag.h\"\n" +
	"TEXT ·f(SB), NOSPLIT, $0\n" +
	"\tMOVQ AX, BX\n" +
	"loop:\n" +
	"\tJMP loop\n" +
	"\tRET\n"

const buggy = "#include \"textflag.h\"\n" +
	"TEXT ·f(SB), NOSPLIT, $0\n" +
	"\tBOGUS AX, BX\n" +
	"\tJMP nowhere\n" +
	"\tRET\n"

func writeTemp(t *testing.T, name, content string) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), name)
	if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
		t.Fatal(err)
	}
	return path
}

// capture runs fn with stdout and stderr redirected and returns both plus the
// exit code fn produced.
func capture(fn func() int) (stdout, stderr string, code int) {
	oldOut, oldErr := os.Stdout, os.Stderr
	rOut, wOut, _ := os.Pipe()
	rErr, wErr, _ := os.Pipe()
	os.Stdout, os.Stderr = wOut, wErr

	code = fn()

	wOut.Close()
	wErr.Close()
	os.Stdout, os.Stderr = oldOut, oldErr
	ob, _ := io.ReadAll(rOut)
	eb, _ := io.ReadAll(rErr)
	return string(ob), string(eb), code
}

// TestCmdFmtRecursive checks the go-fmt-style directory mode: with no
// arguments every .s file below the working directory is formatted in place
// ("." and "_" directories skipped), changed files are listed, and a second
// run is a no-op.
func TestCmdFmtRecursive(t *testing.T) {
	tmp := t.TempDir()
	t.Chdir(tmp)
	unformatted := []byte("TEXT ·f(SB),NOSPLIT,$0\nRET\n")
	write := func(path string) {
		if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(path, unformatted, 0o644); err != nil {
			t.Fatal(err)
		}
	}
	write("a_amd64.s")
	write(filepath.Join("sub", "b_amd64.s"))
	write(filepath.Join("_refs", "c_amd64.s"))
	write(filepath.Join(".git", "d_amd64.s"))

	out, errOut, code := capture(func() int { return cmdFmt(nil) })
	if code != 0 {
		t.Fatalf("code = %d (%s)", code, errOut)
	}
	if out != "a_amd64.s\n"+filepath.Join("sub", "b_amd64.s")+"\n" {
		t.Errorf("listed files unexpected:\n%s", out)
	}
	for _, p := range []string{"a_amd64.s", filepath.Join("sub", "b_amd64.s")} {
		b, _ := os.ReadFile(p)
		if !strings.Contains(string(b), "\tRET") {
			t.Errorf("%s not formatted in place:\n%s", p, b)
		}
	}
	for _, p := range []string{filepath.Join("_refs", "c_amd64.s"), filepath.Join(".git", "d_amd64.s")} {
		b, _ := os.ReadFile(p)
		if string(b) != string(unformatted) {
			t.Errorf("%s must not be touched:\n%s", p, b)
		}
	}

	// Second pass: everything is canonical, nothing is listed.
	out, _, code = capture(func() int { return cmdFmt(nil) })
	if code != 0 || out != "" {
		t.Errorf("second pass: code=%d out=%q, want a no-op", code, out)
	}
}

func TestCmdTokens(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", clean)
	out, _, code := capture(func() int { return cmdTokens([]string{path}) })
	if code != 0 {
		t.Fatalf("code = %d", code)
	}
	if !strings.Contains(out, "IDENT") || !strings.Contains(out, "TEXT") {
		t.Errorf("token dump missing expected tokens:\n%s", out)
	}
}

func TestCmdParseOK(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", clean)
	out, _, code := capture(func() int { return cmdParse([]string{path}) })
	if code != 0 {
		t.Fatalf("code = %d", code)
	}
	if !strings.Contains(out, "OK") || !strings.Contains(out, "1 functions") {
		t.Errorf("parse output = %q", out)
	}
}

func TestCmdParseError(t *testing.T) {
	path := writeTemp(t, "bad_amd64.s", "TEXT ·f(SB), NOSPLIT, $0\n) (\n\tRET\n")
	_, errOut, code := capture(func() int { return cmdParse([]string{path}) })
	if code != 1 {
		t.Fatalf("code = %d, want 1", code)
	}
	if errOut == "" {
		t.Error("expected a parse error on stderr")
	}
}

func TestCmdParseMissingFile(t *testing.T) {
	_, _, code := capture(func() int { return cmdParse([]string{"/nonexistent/file.s"}) })
	if code != 1 {
		t.Fatalf("code = %d, want 1", code)
	}
}

func TestCmdLintClean(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", clean)
	_, _, code := capture(func() int { return cmdLint([]string{path}) })
	if code != 0 {
		t.Fatalf("clean file should lint with code 0, got %d", code)
	}
}

func TestCmdLintErrors(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", buggy)
	out, _, code := capture(func() int { return cmdLint([]string{path}) })
	if code != 1 {
		t.Fatalf("code = %d, want 1", code)
	}
	if !strings.Contains(out, "unknown-instruction") || !strings.Contains(out, "undefined-label") {
		t.Errorf("lint output missing expected codes:\n%s", out)
	}
}

func TestCmdLintDisable(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", buggy)
	args := []string{"-disable", "unknown-instruction,undefined-label", path}
	_, _, code := capture(func() int { return cmdLint(args) })
	if code != 0 {
		t.Fatalf("disabling both rules should yield code 0, got %d", code)
	}
}

func TestCmdFmtStdout(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", "TEXT ·f(SB),NOSPLIT,$0\nMOVQ AX,BX\nRET\n")
	out, _, code := capture(func() int { return cmdFmt([]string{path}) })
	if code != 0 {
		t.Fatalf("code = %d", code)
	}
	if !strings.Contains(out, "TEXT ·f(SB), NOSPLIT, $0") || !strings.Contains(out, "\tMOVQ AX, BX") {
		t.Errorf("formatted output unexpected:\n%s", out)
	}
}

func TestCmdFmtWrite(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", "TEXT ·f(SB),NOSPLIT,$0\nMOVQ AX,BX\nRET\n")
	_, _, code := capture(func() int { return cmdFmt([]string{"-w", path}) })
	if code != 0 {
		t.Fatalf("code = %d", code)
	}
	b, _ := os.ReadFile(path)
	if !strings.Contains(string(b), "\tMOVQ AX, BX") {
		t.Errorf("file not rewritten:\n%s", b)
	}
	// Idempotent: a second -w pass leaves the file unchanged.
	_, _, _ = capture(func() int { return cmdFmt([]string{"-w", path}) })
	b2, _ := os.ReadFile(path)
	if string(b) != string(b2) {
		t.Error("fmt -w is not idempotent")
	}
}

func TestUsage(t *testing.T) {
	var b bytes.Buffer
	usage(&b)
	out := b.String()
	for _, want := range []string{
		"gasm", "Commands:", "Flags:", "--help", "--version",
		"tokens", "parse", "fmt", "lint", "asm", "lsp", "version",
	} {
		if !strings.Contains(out, want) {
			t.Errorf("usage text missing %q:\n%s", want, out)
		}
	}
}

func TestCmdVersion(t *testing.T) {
	out, _, code := capture(func() int { return cmdVersion() })
	if code != 0 {
		t.Fatalf("code = %d", code)
	}
	got := version()
	if !strings.Contains(out, got) {
		t.Errorf("version output %q does not mention %q", out, got)
	}
}

func TestCmdArgErrors(t *testing.T) {
	// A missing path is an error (code 1); cmdFmt with no arguments is the
	// recursive mode now, covered by TestCmdFmtRecursive.
	if _, _, code := capture(func() int { return cmdFmt([]string{"no/such/path"}) }); code != 1 {
		t.Errorf("cmdFmt(missing path) code = %d, want 1", code)
	}
	if _, _, code := capture(func() int { return cmdLint(nil) }); code != 2 {
		t.Errorf("cmdLint() code = %d, want 2", code)
	}
	if _, _, code := capture(func() int { return cmdTokens(nil) }); code != 2 {
		t.Errorf("cmdTokens() code = %d, want 2", code)
	}
	if _, _, code := capture(func() int { return cmdParse(nil) }); code != 2 {
		t.Errorf("cmdParse() code = %d, want 2", code)
	}
}

// TestUsageExitCodes pins the exit-code contract for the commands whose main
// dispatches on a returned error: a wrong argument set exits 2, the same as
// the commands that count their arguments themselves, while a runtime
// failure (an unreadable file) keeps exit 1.
func TestUsageExitCodes(t *testing.T) {
	for name, err := range map[string]error{
		"audit-instructions extra argument": cmdAuditInstructions([]string{"amd64", "extra"}),
		"audit-instructions unknown arch":   cmdAuditInstructions([]string{"mips"}),
		"audit-instructions corpus extra":   cmdAuditInstructions([]string{"--corpus", "a", "b"}),
		"scaffold no arguments":             cmdScaffold(nil),
		"scaffold extra arguments":          cmdScaffold([]string{"differential", "a.s", "b.s"}),
	} {
		if err == nil {
			t.Errorf("%s: expected an error", name)
			continue
		}
		if code := exitCodeFor(err); code != 2 {
			t.Errorf("%s: exit code = %d, want 2 (err: %v)", name, code, err)
		}
	}
	if err := cmdScaffold([]string{"differential", "/nonexistent/file.s"}); err == nil {
		t.Error("scaffold on a missing file should fail")
	} else if code := exitCodeFor(err); code != 1 {
		t.Errorf("scaffold on a missing file: exit code = %d, want 1", code)
	}
}

// TestCmdAsmFormatValidation checks that an unknown --format exits 2 with
// and without -o, instead of assembling and silently dumping a raw image.
func TestCmdAsmFormatValidation(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", clean)
	out := filepath.Join(t.TempDir(), "f.bin")
	if _, _, code := capture(func() int { return cmdAsm([]string{"--format", "bogus", path}) }); code != 2 {
		t.Errorf("asm --format bogus without -o: code = %d, want 2", code)
	}
	if _, _, code := capture(func() int { return cmdAsm([]string{"--format", "bogus", "-o", out, path}) }); code != 2 {
		t.Errorf("asm --format bogus with -o: code = %d, want 2", code)
	}
}

// TestCmdAsmOutputFile pins the documented -o behaviour: the output goes to
// the file and stdout carries no hex dump; without -o the dump is the output.
func TestCmdAsmOutputFile(t *testing.T) {
	path := writeTemp(t, "f_amd64.s", clean)
	out := filepath.Join(t.TempDir(), "f.bin")
	stdout, _, code := capture(func() int { return cmdAsm([]string{"-o", out, path}) })
	if code != 0 {
		t.Fatalf("code = %d", code)
	}
	if strings.Contains(stdout, "0000:") {
		t.Errorf("stdout carries a hex dump despite -o:\n%s", stdout)
	}
	if !strings.Contains(stdout, "wrote ") {
		t.Errorf("stdout misses the wrote line:\n%s", stdout)
	}
	b, err := os.ReadFile(out)
	if err != nil {
		t.Fatal(err)
	}
	if len(b) == 0 {
		t.Error("the output file is empty")
	}

	stdout, _, code = capture(func() int { return cmdAsm([]string{path}) })
	if code != 0 {
		t.Fatalf("without -o: code = %d", code)
	}
	if !strings.Contains(stdout, "0000:") {
		t.Errorf("without -o the hex dump is missing:\n%s", stdout)
	}
}

// TestVerifyNonJITAMD64GroundTruth drives the cross-architecture
// ground-truth path for an amd64 kernel: the path a host of any other
// architecture takes, which must compare against the toolchain rather than
// refuse to run.
func TestVerifyNonJITAMD64GroundTruth(t *testing.T) {
	if testing.Short() {
		t.Skip("runs go tool asm")
	}
	path := writeTemp(t, "f_amd64.s", clean)
	out, _, code := capture(func() int { return cmdVerifyNonJIT(path, arch.AMD64, true, false) })
	if code != 0 {
		t.Fatalf("code = %d (%s)", code, out)
	}
	if !strings.Contains(out, "1/1 matched") {
		t.Errorf("output misses the matched report:\n%s", out)
	}
}

// TestVerifySmokeCrashIsolation checks that a function faulting on its
// zeroed smoke arguments is reported as CRASH by a child process instead of
// killing `gasm verify` itself.
func TestVerifySmokeCrashIsolation(t *testing.T) {
	if testing.Short() {
		t.Skip("builds the gasm binary")
	}
	if runtime.GOARCH != "amd64" {
		t.Skip("amd64 JIT only")
	}
	bin := filepath.Join(t.TempDir(), "gasm")
	if out, err := exec.Command("go", "build", "-o", bin, ".").CombinedOutput(); err != nil {
		t.Fatalf("build gasm: %v\n%s", err, out)
	}
	src := filepath.Join(t.TempDir(), "crash_amd64.s")
	kernel := "#include \"textflag.h\"\n" +
		"\n" +
		"// func Fault(x []byte) int\n" +
		"TEXT ·Fault(SB), NOSPLIT, $0-32\n" +
		"\tMOVQ\tx+0(FP), AX\n" +
		"\tMOVQ\t(AX), AX // faults on the zeroed nil pointer\n" +
		"\tMOVQ\tAX, ret+24(FP)\n" +
		"\tRET\n"
	if err := os.WriteFile(src, []byte(kernel), 0o644); err != nil {
		t.Fatal(err)
	}
	cmd := exec.Command(bin, "verify", "-smoke", src)
	out, err := cmd.CombinedOutput()
	if err == nil {
		t.Fatalf("expected a failure report, got success:\n%s", out)
	}
	if exitErr, ok := err.(*exec.ExitError); ok {
		if ws, ok := exitErr.Sys().(syscall.WaitStatus); ok && ws.Signaled() {
			t.Fatalf("verify died from %v; the crash was not isolated:\n%s", ws.Signal(), out)
		}
	}
	if !strings.Contains(string(out), "CRASH") {
		t.Errorf("output does not report CRASH:\n%s", out)
	}
}

func TestSweepCheckLines(t *testing.T) {
	out := []byte("crash_amd64.s: 1 functions JIT-loaded\n" +
		"  Fault: 21 bytes, args=32, frame=0 NOSPLIT\n" +
		"    smoke: OK\n" +
		"    abi: clean (10 varied inputs)\n")
	want := "    smoke: OK\n    abi: clean (10 varied inputs)"
	if got := sweepCheckLines(out); got != want {
		t.Errorf("sweepCheckLines = %q, want %q", got, want)
	}
}

// TestRunCorpusAudit drives the corpus audit over a small fixture tree: one
// suffixed amd64 file, one suffixed arm64 file whose body is not arm64, one
// generic file, and one file that does not parse.
func TestRunCorpusAudit(t *testing.T) {
	dir := t.TempDir()
	write := func(name, src string) {
		t.Helper()
		if err := os.WriteFile(filepath.Join(dir, name), []byte(src), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	write("good_amd64.s", "#include \"textflag.h\"\nTEXT ·add(SB), NOSPLIT, $0-0\n\tMOVQ AX, BX\n\tRET\n")
	write("bad_arm64.s", "#include \"textflag.h\"\nTEXT ·f(SB), NOSPLIT, $0-0\n\tMOVQ AX, BX\n\tRET\n")
	write("generic.s", "#include \"textflag.h\"\nTEXT ·g(SB), NOSPLIT, $0-0\n\tRET\n")
	write("broken.s", "#include \"textflag.h\"\nTEXT ·b(SB), NOSPLIT, $0-0\n\tJMP nowhere\n\tRET\n")

	stats, err := runCorpusAudit(dir)
	if err != nil {
		t.Fatalf("runCorpusAudit: %v", err)
	}
	if stats.files != 4 {
		t.Errorf("files = %d, want 4", stats.files)
	}
	if stats.generic != 2 {
		t.Errorf("generic = %d, want 2 (generic.s and broken.s)", stats.generic)
	}
	// good_amd64 and generic.s assemble everywhere they are attempted.
	if stats.full != 2 {
		t.Errorf("full = %d, want 2", stats.full)
	}
	get := func(name string) *corpusTally {
		for i, tg := range stats.targets {
			if tg.name == name {
				return stats.tallies[i]
			}
		}
		t.Fatalf("no tally for %s", name)
		return nil
	}
	// amd64: good_amd64 + generic.s + broken.s; the broken file fails to parse.
	if a := get("amd64"); a.attempted != 3 || a.assembled != 2 {
		t.Errorf("amd64 = %d/%d, want 2/3", a.assembled, a.attempted)
	}
	// arm64: bad_arm64 (MOVQ is not arm64) + generic.s + broken.s.
	if a := get("arm64"); a.attempted != 3 || a.assembled != 1 {
		t.Errorf("arm64 = %d/%d, want 1/3", a.assembled, a.attempted)
	}
	if r := get("amd64").reasons["instruction not encodable"]; r != 0 {
		t.Errorf("amd64 unexpected unencodable reason: %d", r)
	}
	if r := get("arm64").reasons["instruction not encodable"]; r != 1 {
		t.Errorf("arm64 unencodable reasons = %d, want 1", r)
	}
}