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

package parser

import (
	"os"
	"path/filepath"
	"strings"
	"testing"

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

// expand parses src with preprocessing enabled and returns the first TEXT's
// body instructions as "MNEMONIC operand|operand" strings, the shape the
// expansion assertions below compare against.  Runs of spaces are
// collapsed: Raw renders a token group as its tokens joined with single
// spaces, so "$(32-7)" arrives as "$ ( 32 - 7 )" and the comparison must
// not depend on that spelling.
func expand(t *testing.T, src string) (*ast.File, []string) {
	t.Helper()
	f, errs := ParseWithOptions("t_amd64.s", src, Options{Expand: true})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	ts := texts(f)
	if len(ts) == 0 {
		t.Fatalf("no TEXT in:\n%s", src)
	}
	var got []string
	for _, s := range ts[0].Body {
		in, ok := s.(*ast.Instr)
		if !ok {
			continue
		}
		var ops []string
		for _, op := range in.Operands {
			ops = append(ops, op.Raw)
		}
		line := in.Mnemonic.Text + " " + strings.Join(ops, ", ")
		got = append(got, strings.ReplaceAll(line, " ", ""))
	}
	return f, got
}

func wantLines(t *testing.T, got []string, want ...string) {
	t.Helper()
	strip := func(lines []string) string {
		var out []string
		for _, l := range lines {
			out = append(out, strings.ReplaceAll(l, " ", ""))
		}
		return strings.Join(out, "\n")
	}
	if strip(got) != strip(want) {
		t.Errorf("expanded body:\n  %s\nwant:\n  %s", strings.Join(got, "\n  "), strings.Join(want, "\n  "))
	}
}

func TestObjectMacroExpandsAtUse(t *testing.T) {
	_, got := expand(t, `
#define REGTMP CX
#define TWICE ADDQ CX, AX; ADDQ CX, AX
TEXT ·f(SB), NOSPLIT, $0
	MOVQ 8(SP), REGTMP
	TWICE
	RET
`)
	wantLines(t, got,
		"MOVQ 8(SP), CX",
		"ADDQ CX, AX",
		"ADDQ CX, AX",
		"RET",
	)
}

func TestParameterisedMacroSubstitutesArguments(t *testing.T) {
	f, errs := ParseWithOptions("t_amd64.s", `
#define ROUND1(a, index, const, shift) \
	ADDQ $const, a; \
	MOVW (index*4)(SP), a; \
	RORQ $(32-shift), a
TEXT ·f(SB), NOSPLIT, $0
	ROUND1(AX, 3, 0xd76aa478, 7)
	RET
`, Options{Expand: true})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	body := texts(f)[0].Body
	add := body[0].(*ast.Instr)
	if add.Mnemonic.Text != "ADDQ" || !add.Operands[0].Imm.HasVal ||
		add.Operands[0].Imm.Val != 0xd76aa478 || add.Operands[1].Addr.Sym == nil ||
		add.Operands[1].Addr.Sym.Name != "AX" {
		t.Errorf("ADDQ operands substituted wrong: %+v %+v", add.Operands[0].Imm, add.Operands[1].Addr)
	}
	mov := body[1].(*ast.Instr)
	if addr := mov.Operands[0].Addr; !addr.HasOff || addr.Offset != 12 {
		t.Errorf("MOVW offset = %+v, want 12 from 3*4", addr)
	}
	ror := body[2].(*ast.Instr)
	if !ror.Operands[0].Imm.HasVal || ror.Operands[0].Imm.Val != 25 {
		t.Errorf("RORQ immediate = %+v, want 25 from (32-7)", ror.Operands[0].Imm)
	}
}

func TestMacroArgumentsKeepCommasInParens(t *testing.T) {
	// An argument may itself be an unparenthesised expression: the tokens
	// substitute verbatim and the parser folds the result, as the
	// toolchain's parser does.
	f, errs := ParseWithOptions("t_amd64.s", `
#define LOAD(dst, off) MOVQ off(SP), dst
TEXT ·f(SB), NOSPLIT, $0
	LOAD(AX, 1*8)
	RET
`, Options{Expand: true})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	in := texts(f)[0].Body[0].(*ast.Instr)
	addr := in.Operands[0].Addr
	if !addr.HasOff || addr.Offset != 8 {
		t.Errorf("offset = %+v, want 8", addr)
	}
	if sym := in.Operands[1].Addr.Sym; sym == nil || sym.Name != "AX" {
		t.Errorf("destination = %+v, want AX", in.Operands[1].Addr)
	}
}

func TestNestedMacroInvocations(t *testing.T) {
	// An object macro naming a parameterised one, and a parameterised body
	// invoking another parameterised macro: the toolchain's input stack
	// rescans substituted tokens, and so does expansion here.
	_, got := expand(t, `
#define DOUBLE(x) ADDQ x, x
#define TWICE2 DOUBLE
#define FOUR(a, b) DOUBLE(a); DOUBLE(b)
TEXT ·f(SB), NOSPLIT, $0
	TWICE2(AX)
	FOUR(AX, CX)
	RET
`)
	wantLines(t, got,
		"ADDQ AX, AX",
		"ADDQ AX, AX",
		"ADDQ CX, CX",
		"RET",
	)
}

func TestMultiLineBodySplitsWithoutSemicolons(t *testing.T) {
	// The arm64 style: backslash-continued lines with no semicolons.  The
	// continuation newline is a statement boundary, as in the toolchain.
	_, got := expand(t, `
#define PAIR \
	ADDQ AX, AX \
	MOVQ AX, CX
TEXT ·f(SB), NOSPLIT, $0
	PAIR
	RET
`)
	wantLines(t, got,
		"ADDQ AX, AX",
		"MOVQ AX, CX",
		"RET",
	)
}

func TestZeroArgumentMacro(t *testing.T) {
	_, got := expand(t, `
#define BARRIER()
TEXT ·f(SB), NOSPLIT, $0
	BARRIER()
	RET
`)
	wantLines(t, got, "RET")
}

func TestParameterisedWithoutParensStandsAsName(t *testing.T) {
	// A parameterised macro invoked without its parentheses names itself,
	// which the parser then reports as an unknown instruction rather than
	// silently expanding nothing.
	f, errs := ParseWithOptions("t_amd64.s", `
#define M(x) ADDQ x, x
TEXT ·f(SB), NOSPLIT, $0
	M
	RET
`, Options{Expand: true})
	if len(errs) != 0 {
		t.Fatalf("parse: %v", errs)
	}
	fn := texts(f)[0]
	if len(fn.Body) == 0 {
		t.Fatal("body empty")
	}
	in, ok := fn.Body[0].(*ast.Instr)
	if !ok || in.Mnemonic.Text != "M" {
		t.Fatalf("bare parameterised macro did not stand as its name: %+v", fn.Body[0])
	}
}

func TestDefinitionScoping(t *testing.T) {
	// A definition applies from its point onward: the use before the
	// #define stays untouched.
	_, got := expand(t, `
TEXT ·f(SB), NOSPLIT, $0
	SPECIAL
#define SPECIAL ADDQ AX, AX
	SPECIAL
	RET
`)
	wantLines(t, got,
		"SPECIAL",
		"ADDQ AX, AX",
		"RET",
	)
}

func TestUndefRemovesMacro(t *testing.T) {
	_, got := expand(t, `
#define TEMP AX
TEXT ·f(SB), NOSPLIT, $0
	TEMP
#undef TEMP
	TEMP
	RET
`)
	wantLines(t, got,
		"AX",
		"TEMP",
		"RET",
	)
}

func TestUndefUndefinedMacroIsAnError(t *testing.T) {
	_, errs := ParseWithOptions("t_amd64.s", "#undef NOSUCH\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n", Options{Expand: true})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "undefined macro NOSUCH") {
		t.Fatalf("#undef of an undefined macro: got %v, want an error naming it", errs)
	}
}

func TestRedefinitionIsAnError(t *testing.T) {
	_, errs := ParseWithOptions("t_amd64.s", "#define A X\n#define A Y\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n", Options{Expand: true})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "redefinition of macro A") {
		t.Fatalf("redefinition: got %v, want an error", errs)
	}
}

func TestRecursiveMacroIsAnError(t *testing.T) {
	_, errs := ParseWithOptions("t_amd64.s", "#define A B\n#define B A\nTEXT ·f(SB), NOSPLIT, $0\n\tA\n\tRET\n", Options{Expand: true})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "recursive macro invocation") {
		t.Fatalf("recursion: got %v, want a recursive-macro error, not a hang", errs)
	}
}

func TestWrongArgumentCountIsAnError(t *testing.T) {
	_, errs := ParseWithOptions("t_amd64.s", "#define M(a, b) ADDQ a, b\nTEXT ·f(SB), NOSPLIT, $0\n\tM(AX)\n\tRET\n", Options{Expand: true})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "wrong arg count for macro M") {
		t.Fatalf("arg count: got %v, want an error", errs)
	}
}

func TestConditionalsSelectOneBranch(t *testing.T) {
	_, got := expand(t, `
#define MODE2
TEXT ·f(SB), NOSPLIT, $0
#ifdef MODE2
	ADDQ AX, AX
#else
	SUBQ AX, AX
#endif
#ifndef MODE2
	SUBQ CX, CX
#else
	ADDQ CX, CX
#endif
	RET
`)
	wantLines(t, got,
		"ADDQ AX, AX",
		"ADDQ CX, CX",
		"RET",
	)
}

func TestConditionalsHideDefinitionsAndIncludes(t *testing.T) {
	// A definition inside a disabled branch must not exist, and an
	// unresolvable include there must not be followed.
	_, got := expand(t, `
TEXT ·f(SB), NOSPLIT, $0
#ifdef NOTDEFINED
#define HIDEN ADDQ AX, AX
#include "nowhere.h"
#endif
	HIDEN
	RET
`)
	wantLines(t, got, "HIDEN", "RET")
}

func TestUnclosedConditionalIsAnError(t *testing.T) {
	_, errs := ParseWithOptions("t_amd64.s", "#ifdef X\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n", Options{Expand: true})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "unclosed #ifdef") {
		t.Fatalf("unclosed conditional: got %v, want an error", errs)
	}
}

func TestUnmatchedConditionalDelimitersAreErrors(t *testing.T) {
	_, errs := ParseWithOptions("t_amd64.s", "#endif\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n", Options{Expand: true})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "unmatched #endif") {
		t.Fatalf("unmatched #endif: got %v, want an error", errs)
	}
	_, errs = ParseWithOptions("t_amd64.s", "#else\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n", Options{Expand: true})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "unmatched #else") {
		t.Fatalf("unmatched #else: got %v, want an error", errs)
	}
}

// includeTree writes a directory of include files and returns its path.
func includeTree(t *testing.T, files map[string]string) string {
	t.Helper()
	dir := t.TempDir()
	for name, content := range files {
		path := filepath.Join(dir, name)
		if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil {
			t.Fatal(err)
		}
		if err := os.WriteFile(path, []byte(content), 0o644); err != nil {
			t.Fatal(err)
		}
	}
	return dir
}

func TestIncludeSplicesAndDefinesAreShared(t *testing.T) {
	dir := includeTree(t, map[string]string{
		"consts.h": "#define KONST $42\n",
	})
	f, errs := ParseWithOptions("t_amd64.s", `
#include "consts.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVQ KONST, AX
	RET
`, Options{Expand: true, IncludeDirs: []string{dir}})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	in := texts(f)[0].Body[0].(*ast.Instr)
	if in.Mnemonic.Text != "MOVQ" || strings.ReplaceAll(in.Operands[0].Raw, " ", "") != "$42" {
		t.Fatalf("include splicing failed: %+v", in)
	}
}

func TestIncludeResolutionOrder(t *testing.T) {
	// The including file's directory wins over the -I list, and the -I list
	// is searched in order.
	src := includeTree(t, map[string]string{
		"inc/main.s":     "#include \"which.h\"\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n",
		"inc/which.h":    "#define WHO ONE\n",
		"first/which.h":  "#define WHO TWO\n",
		"second/which.h": "#define WHO THREE\n",
	})
	main := filepath.Join(src, "inc", "main.s")
	body, err := os.ReadFile(main)
	if err != nil {
		t.Fatal(err)
	}
	// The header exists in the including file's directory and in two -I
	// directories; the source-directory copy must win.
	f, errs := ParseWithOptions(main, string(body), Options{Expand: true, IncludeDirs: []string{
		filepath.Join(src, "first"), filepath.Join(src, "second"),
	}})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	found := false
	for _, d := range f.Decls {
		if pp, ok := d.(*ast.Preproc); ok && strings.Contains(pp.Raw, "define WHO ONE") {
			found = true
		}
	}
	if !found {
		t.Error("the including file's directory did not win include resolution")
	}
}

func TestIncludeSearchesIncludeDirsInOrder(t *testing.T) {
	src := includeTree(t, map[string]string{
		"inc/main.s":     "#include \"which.h\"\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n",
		"first/which.h":  "#define WHO TWO\n",
		"second/which.h": "#define WHO THREE\n",
	})
	main := filepath.Join(src, "inc", "main.s")
	body, err := os.ReadFile(main)
	if err != nil {
		t.Fatal(err)
	}
	f, errs := ParseWithOptions(main, string(body), Options{Expand: true, IncludeDirs: []string{
		filepath.Join(src, "first"), filepath.Join(src, "second"),
	}})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	for _, d := range f.Decls {
		if pp, ok := d.(*ast.Preproc); ok && strings.Contains(pp.Raw, "define WHO THREE") {
			t.Error("the second -I directory was searched before the first")
		}
	}
}

func TestIncludeCycleIsDetected(t *testing.T) {
	src := includeTree(t, map[string]string{
		"a.s": "#include \"b.s\"\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n",
		"b.s": "#include \"a.s\"\n",
	})
	_, errs := ParseWithOptions(filepath.Join(src, "a.s"), "#include \"b.s\"\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n",
		Options{Expand: true})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), "include cycle") {
		t.Fatalf("include cycle: got %v, want a cycle diagnostic, not a hang", errs)
	}
}

func TestUnresolvableIncludeIsAnError(t *testing.T) {
	_, errs := ParseWithOptions("t_amd64.s", "#include \"nothere.h\"\nTEXT ·f(SB), NOSPLIT, $0\n\tRET\n",
		Options{Expand: true, IncludeDirs: []string{t.TempDir()}})
	if len(errs) == 0 || !strings.Contains(errs[0].Error(), `#include "nothere.h"`) {
		t.Fatalf("missing include: got %v, want a clear diagnostic", errs)
	}
}

func TestTextflagHeaderIsNeverSpliced(t *testing.T) {
	// textflag.h resolves nowhere here, yet the file must parse: the flag
	// names are consumed natively and the include stays in the tree.
	f, errs := ParseWithOptions("t_amd64.s", `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	RET
`, Options{Expand: true})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	hasInclude := false
	for _, d := range f.Decls {
		if _, ok := d.(*ast.Include); ok {
			hasInclude = true
		}
	}
	if !hasInclude {
		t.Error("textflag.h include was dropped from the tree")
	}
}

func TestSemicolonSplitsRawLinesToo(t *testing.T) {
	_, got := expand(t, `
TEXT ·f(SB), NOSPLIT, $0
	BYTE $0x0f; BYTE $0x1f
	RET
`)
	wantLines(t, got, "BYTE $0x0f", "BYTE $0x1f", "RET")
}

func TestParseUnchangedWithoutExpand(t *testing.T) {
	// Without Expand the preprocessor must not exist: a macro invocation
	// stays an unexpanded instruction line and ';' keeps the old parse.
	f, errs := Parse("t_amd64.s", `
#define TWICE ADDQ AX, AX
TEXT ·f(SB), NOSPLIT, $0
	TWICE
	BYTE $0x0f; BYTE $0x1f
	RET
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	fn := texts(f)[0]
	var mnemonics []string
	for _, s := range fn.Body {
		if in, ok := s.(*ast.Instr); ok {
			mnemonics = append(mnemonics, in.Mnemonic.Text)
		}
	}
	if strings.Join(mnemonics, " ") != "TWICE BYTE RET" {
		t.Errorf("non-expanding parse changed: %v", mnemonics)
	}
}

func TestConstantExpressionFolding(t *testing.T) {
	// The shapes substituted macro bodies leave behind: parenthesised
	// arithmetic in immediates and displacements, tilde complements.  The
	// assertions read the semantic fields; Raw keeps the operand's tokens
	// in the canonicalised rendering, not the folded values.
	f, errs := ParseWithOptions("t_amd64.s", `
TEXT ·f(SB), NOSPLIT, $0
	RORQ $(32-7), AX
	ANDQ $~63, AX
	MOVQ ((2*4)+0)(SP), AX
	MOVQ $((1<<3)|(1<<1)), AX
	RET
`, Options{Expand: true})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	body := texts(f)[0].Body
	ror := body[0].(*ast.Instr)
	if !ror.Operands[0].Imm.HasVal || ror.Operands[0].Imm.Val != 25 {
		t.Errorf("RORQ immediate = %+v, want 25", ror.Operands[0].Imm)
	}
	and := body[1].(*ast.Instr)
	if !and.Operands[0].Imm.HasVal || and.Operands[0].Imm.Val != -64 {
		t.Errorf("ANDQ immediate = %+v, want -64", and.Operands[0].Imm)
	}
	mov := body[2].(*ast.Instr)
	addr := mov.Operands[0].Addr
	if !addr.HasOff || addr.Offset != 8 || addr.Base != "SP" {
		t.Errorf("MOVQ address = %+v, want 8(SP)", addr)
	}
	mov2 := body[3].(*ast.Instr)
	if !mov2.Operands[0].Imm.HasVal || mov2.Operands[0].Imm.Val != 10 {
		t.Errorf("MOVQ immediate = %+v, want 10", mov2.Operands[0].Imm)
	}
}

func TestConstantExpressionFoldsWithoutExpand(t *testing.T) {
	// Folding is a parser capability, not a preprocessing one: a
	// hand-written $(32-7) folds the same way with expansion off.
	f, errs := ParseWithOptions("t_amd64.s", "TEXT ·f(SB), NOSPLIT, $0\n\tRORQ $(32-7), AX\n\tRET\n", Options{})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	in := texts(f)[0].Body[0].(*ast.Instr)
	if !in.Operands[0].Imm.HasVal || in.Operands[0].Imm.Val != 25 {
		t.Errorf("Imm = %+v, want 25", in.Operands[0].Imm)
	}
}

func TestNotAnExpressionFallsBack(t *testing.T) {
	// Symbol immediates and floats must keep their ordinary parse.
	f, errs := ParseWithOptions("t_amd64.s", "TEXT ·f(SB), NOSPLIT, $0\n\tMOVQ $1.5, AX\n\tMOVQ $·sym(SB), AX\n\tRET\n", Options{})
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	fn := texts(f)[0]
	mov1 := fn.Body[0].(*ast.Instr)
	if mov1.Operands[0].Imm.HasVal || mov1.Operands[0].Imm.Float != "1.5" {
		t.Errorf("float immediate parsed as %+v", mov1.Operands[0].Imm)
	}
	mov2 := fn.Body[1].(*ast.Instr)
	if mov2.Operands[0].Imm.Sym == nil {
		t.Errorf("symbol immediate parsed as %+v", mov2.Operands[0].Imm)
	}
}