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

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

func TestUsedLabelNotFlagged(t *testing.T) {
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
loop:
	ADDQ $1, AX
	JMP loop
	RET
`)
	if codes(diags)[CodeUnusedLabel] != 0 {
		t.Fatalf("used label must not be flagged: %+v", diags)
	}
}

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

func TestValidTextflags(t *testing.T) {
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT|NOFRAME|DUPOK, $0
	RET
`)
	if codes(diags)[CodeInvalidFlag] != 0 {
		t.Fatalf("valid flags must not be flagged: %+v", diags)
	}
}

func TestStackImbalance(t *testing.T) {
	// Function with frame size 16 but only SUB 8, SP — imbalance.
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $16-0
	SUBQ $8, SP
	RET
`)
	if codes(diags)[CodeStackImbalance] != 1 {
		t.Fatalf("want one stack-imbalance, got %+v", diags)
	}
}

func TestStackBalanced(t *testing.T) {
	// Function with frame size 16 and matching SUB/ADD — balanced.
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $16-0
	SUBQ $16, SP
	ADDQ $16, SP
	RET
`)
	if codes(diags)[CodeStackImbalance] != 0 {
		t.Fatalf("balanced stack must not be flagged: %+v", diags)
	}
}

func TestRegisterWidthMismatch(t *testing.T) {
	// MOVQ with 32-bit register — mismatch.
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVQ EAX, BX
	RET
`)
	if codes(diags)[CodeRegisterWidthMismatch] != 1 {
		t.Fatalf("want one register-width-mismatch, got %+v", diags)
	}
}

func TestRegisterWidthCorrect(t *testing.T) {
	// MOVQ with 64-bit registers — correct.
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVQ RAX, RBX
	RET
`)
	if codes(diags)[CodeRegisterWidthMismatch] != 0 {
		t.Fatalf("correct width must not be flagged: %+v", diags)
	}
}

func TestABI0RegisterArgs(t *testing.T) {
	// A kernel with a // func signature whose parameters are never read from
	// the FP frame: the classic register-args port bug.
	diags := lintSrc(t, `
// func kernel(text *byte, n int)
TEXT ·kernel(SB), NOSPLIT, $0-16
	MOVQ DI, R10
	MOVQ R9, AX
	RET
`)
	if codes(diags)[CodeABI0RegisterArgs] != 1 {
		t.Fatalf("want one abi0-register-args, got %+v", diags)
	}

	// Reading every parameter from the frame is correct.
	diags = lintSrc(t, `
// func kernel(text *byte, n int)
TEXT ·kernel(SB), NOSPLIT, $0-16
	MOVQ text+0(FP), AX
	MOVQ n+8(FP), BX
	RET
`)
	if codes(diags)[CodeABI0RegisterArgs] != 0 {
		t.Fatalf("frame-reading kernel must not be flagged: %+v", diags)
	}

	// Only a result write to FP, no parameter read: still flagged.
	diags = lintSrc(t, `
// func top(sa []int32) int
TEXT ·top(SB), NOSPLIT, $0-24
	MOVQ SI, R10
	MOVQ R10, ret+16(FP)
	RET
`)
	if codes(diags)[CodeABI0RegisterArgs] != 1 {
		t.Fatalf("result-only FP write must still be flagged: %+v", diags)
	}

	// No signature: stay silent.
	diags = lintSrc(t, `
TEXT ·bare(SB), NOSPLIT, $0-16
	MOVQ DI, AX
	RET
`)
	if codes(diags)[CodeABI0RegisterArgs] != 0 {
		t.Fatalf("signature-less function must not be flagged: %+v", diags)
	}
}

func TestRegisterWidthCanonicalNames(t *testing.T) {
	// Canonical Go asm names with an L operation: the correct spelling for a
	// 32-bit operation, never a width mismatch.
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVL (BX)(R9*4), SI
	XORL R9, R9
	MOVL 128(BX)(R9*4), R12
	RET
`)
	if codes(diags)[CodeRegisterWidthMismatch] != 0 {
		t.Fatalf("canonical 64-bit names with L ops must not be flagged: %+v", diags)
	}

	// A byte register in an L operation stays a mismatch.
	diags = lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVL AL, DX
	RET
`)
	if codes(diags)[CodeRegisterWidthMismatch] != 1 {
		t.Fatalf("byte register in L op must be flagged: %+v", diags)
	}
}

func TestNonportableRegisterName(t *testing.T) {
	diags := lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVQ RAX, RBX
	MOVQ EAX, R9
	RET
`)
	// RAX and RBX (and EAX) are gasm aliases of AX/BX; R9 is canonical.
	if got := codes(diags)[CodeNonportableRegister]; got != 3 {
		t.Fatalf("want 3 nonportable-register-name, got %d: %+v", got, diags)
	}

	diags = lintSrc(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVQ AX, BX
	MOVQ (BX)(R9*4), R12
	RET
`)
	if codes(diags)[CodeNonportableRegister] != 0 {
		t.Fatalf("canonical names must not be flagged: %+v", diags)
	}
}