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

package lint

import (
	"os"
	"testing"

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

func lintSrc(t *testing.T, src string) []Diagnostic {
	t.Helper()
	f, errs := parser.Parse("test_amd64.s", src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	return File(f, Config{Arch: arch.AMD64})
}

// lintSrcArch lints src under the architecture inferred from filename.
func lintSrcArch(t *testing.T, filename, src string) []Diagnostic {
	t.Helper()
	f, errs := parser.Parse(filename, src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	return File(f, Config{Arch: arch.FromFilename(filename)})
}

func codes(diags []Diagnostic) map[string]int {
	m := map[string]int{}
	for _, d := range diags {
		m[d.Code]++
	}
	return m
}

func TestFixtureIsClean(t *testing.T) {
	src, err := os.ReadFile("../testdata/sample_amd64.s")
	if err != nil {
		t.Fatal(err)
	}
	f, errs := parser.Parse("sample_amd64.s", string(src))
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	// The fixture mirrors the go-flac kernels, which write the Go ABI0
	// scratch registers (BX, R13) without saving them — legal under Go's
	// stack-based ABI, so the register-clobber audit stays silent and the
	// fixture must lint entirely clean.
	diags := File(f, Config{Arch: arch.AMD64})
	if len(diags) != 0 {
		t.Fatalf("expected no diagnostics on the fixture, got %+v", diags)
	}
}

func TestUnknownInstruction(t *testing.T) {
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	FOOBAR AX, BX
	RET
`)
	if codes(diags)[CodeUnknownInstr] != 1 {
		t.Fatalf("want one unknown-instruction, got %+v", diags)
	}
}

func TestUndefinedLabel(t *testing.T) {
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	JMP nowhere
	RET
`)
	if codes(diags)[CodeUndefinedLabel] != 1 {
		t.Fatalf("want one undefined-label, got %+v", diags)
	}
}

func TestDuplicateLabel(t *testing.T) {
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
loop:
	ADDQ $1, AX
loop:
	SUBQ $1, AX
	JMP loop
	RET
`)
	if codes(diags)[CodeDuplicateLabel] != 1 {
		t.Fatalf("want one duplicate-label, got %+v", diags)
	}
}

func TestMissingRet(t *testing.T) {
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	ADDQ $1, AX
`)
	if codes(diags)[CodeMissingRet] != 1 {
		t.Fatalf("want one missing-ret, got %+v", diags)
	}
}

func TestOperandCount(t *testing.T) {
	// RET takes zero operands; JMP takes exactly one.
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	RET AX
	JMP
	RET
`)
	c := codes(diags)
	if c[CodeOperandCount] != 2 {
		t.Fatalf("want two operand-count findings, got %+v", diags)
	}
}

func TestMissingTextflag(t *testing.T) {
	diags := lintSrc(t, `
TEXT ·f(SB), NOSPLIT, $0
	RET
`)
	if codes(diags)[CodeMissingTextflag] != 1 {
		t.Fatalf("want one missing-textflag-include, got %+v", diags)
	}
}

func TestDisableRule(t *testing.T) {
	f, _ := parser.Parse("t_amd64.s", `
TEXT ·f(SB), NOSPLIT, $0
	RET
`)
	diags := File(f, Config{Arch: arch.AMD64, Disable: map[string]bool{CodeMissingTextflag: true}})
	if len(diags) != 0 {
		t.Fatalf("disabling the rule should silence it, got %+v", diags)
	}
}

func TestMacroInvocationSkipped(t *testing.T) {
	// DISPATCH is defined in-file; get_tls carries an underscore.  Neither is a
	// machine instruction, so both must be ignored by the unknown-instruction
	// rule rather than flagged.
	diags := lintSrc(t, `
#include "textflag.h"
#define DISPATCH CALL ·x(SB)
TEXT ·f(SB), NOSPLIT, $0
	DISPATCH
	get_tls CX
	RET
`)
	if codes(diags)[CodeUnknownInstr] != 0 {
		t.Fatalf("macro invocations must not be flagged: %+v", diags)
	}
}

func TestUndefIsTerminal(t *testing.T) {
	// A function whose body is UNDEF traps and never returns; it needs no RET.
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	UNDEF
`)
	if codes(diags)[CodeMissingRet] != 0 {
		t.Fatalf("UNDEF should count as terminal: %+v", diags)
	}
}

func TestArm64BranchAlias(t *testing.T) {
	diags := lintSrcArch(t, "f_arm64.s", `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	B done
done:
	RET
`)
	if len(diags) != 0 {
		t.Fatalf("arm64 B to a defined label should be clean: %+v", diags)
	}
}

// TestEvexMaskingRecognised checks that masked EVEX forms — the .Z suffix and
// an explicit K operand — are recognised and exempt from operand-count
// checks.
func TestEvexMaskingRecognised(t *testing.T) {
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	VPADDD.Z Z1, Z2, K2, Z3
	VPMINSD Z1, Z2, K5, Z3
	VMOVDQU8 Z1, K3, (SI)
	RET
`)
	if codes(diags)[CodeUnknownInstr] != 0 {
		t.Fatalf("masked EVEX must be recognised: %+v", diags)
	}
	if codes(diags)[CodeOperandCount] != 0 {
		t.Fatalf("masked operand counts must not be flagged: %+v", diags)
	}
}

func TestArm64AddressingSuffix(t *testing.T) {
	// .W (pre-index) and .P (post-index) suffixes must resolve to the base
	// instruction.
	diags := lintSrcArch(t, "f_arm64.s", `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	LDP.W (R0), (R1, R2)
	VST1.P (R3), (R4)
	RET
`)
	if codes(diags)[CodeUnknownInstr] != 0 {
		t.Fatalf("suffixed load/store should be recognised: %+v", diags)
	}
}

func TestMacrosInPlaySuppressesLabelRules(t *testing.T) {
	// Including a non-textflag header means macros may define labels and supply
	// the RET, so undefined-label and missing-ret are suppressed.
	diags := lintSrcArch(t, "f_arm64.s", `
#include "go_asm.h"
TEXT ·f(SB), NOSPLIT, $0
	JMP RARG0
`)
	if codes(diags)[CodeUndefinedLabel] != 0 || codes(diags)[CodeMissingRet] != 0 {
		t.Fatalf("label rules should be suppressed in macro files: %+v", diags)
	}
}