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

package asm

import (
	"bytes"
	"strings"
	"testing"

	"golang.org/x/arch/x86/x86asm"

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

// firstText parses src and returns its first TEXT function.
func firstText(t *testing.T, src string) *ast.Text {
	t.Helper()
	f, errs := parser.Parse("f_amd64.s", src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	for _, d := range f.Decls {
		if txt, ok := d.(*ast.Text); ok {
			return txt
		}
	}
	t.Fatal("no TEXT function found")
	return nil
}

// disasm decodes a machine-code blob into Intel-syntax instruction strings.
func disasm(t *testing.T, code []byte) []string {
	t.Helper()
	var out []string
	for len(code) > 0 {
		inst, err := x86asm.Decode(code, 64)
		if err != nil {
			t.Fatalf("decode %x: %v", code, err)
		}
		out = append(out, x86asm.IntelSyntax(inst, 0, nil))
		code = code[inst.Len:]
	}
	return out
}

func hexBytes(b []byte) string {
	var sb strings.Builder
	for _, x := range b {
		sb.WriteString(" ")
		const hexdig = "0123456789abcdef"
		sb.WriteByte(hexdig[x>>4])
		sb.WriteByte(hexdig[x&0xf])
	}
	return strings.TrimSpace(sb.String())
}

func TestAssembleLoop(t *testing.T) {
	fn := firstText(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	XORQ AX, AX
loop:
	ADDQ $1, AX
	CMPQ AX, $10
	JLT loop
	RET
`)
	code, labels, err := Assemble(fn)
	if err != nil {
		t.Fatalf("Assemble: %v", err)
	}
	if _, ok := labels["loop"]; !ok {
		t.Fatalf("label 'loop' not recorded: %v", labels)
	}
	got := strings.Join(disasm(t, code), "\n")
	want := strings.Join([]string{
		"xor rax, rax",
		"add rax, 0x1",
		"cmp rax, 0xa",
		"jl 0x0",
		"ret",
	}, "\n")
	gotLines := strings.Split(got, "\n")
	wantLines := strings.Split(want, "\n")
	if len(gotLines) != len(wantLines) {
		t.Fatalf("instruction count mismatch:\n got:\n%s\n want:\n%s", got, want)
	}
	for i := range wantLines {
		if strings.HasPrefix(wantLines[i], "jl") {
			if !strings.HasPrefix(gotLines[i], "jl") {
				t.Errorf("line %d: got %q, want a jl", i, gotLines[i])
			}
			continue
		}
		if gotLines[i] != wantLines[i] {
			t.Errorf("line %d: got %q, want %q", i, gotLines[i], wantLines[i])
		}
	}
}

func TestAssembleMemory(t *testing.T) {
	fn := firstText(t, `
#include "textflag.h"
TEXT ·g(SB), NOSPLIT, $0
	MOVQ (AX), BX
	MOVQ 8(AX), CX
	LEAQ (AX)(BX*4), DX
	RET
`)
	code, _, err := Assemble(fn)
	if err != nil {
		t.Fatalf("Assemble: %v", err)
	}
	got := strings.Join(disasm(t, code), "\n")
	want := strings.Join([]string{
		"mov rbx, qword ptr [rax]",
		"mov rcx, qword ptr [rax+0x8]",
		"lea rdx, ptr [rax+4*rbx]",
		"ret",
	}, "\n")
	if got != want {
		t.Errorf("assemble memory:\n got:\n%s\n want:\n%s", got, want)
	}
}

// TestAssembleFP verifies the FP pseudo-register translation for a NOSPLIT $0
// function against the exact bytes the Go assembler produces (verified via
// `go tool objdump`): x+N(FP) maps to (N+8)(SP).
func TestAssembleFP(t *testing.T) {
	fn := firstText(t, `
#include "textflag.h"
TEXT ·loadarg(SB), NOSPLIT, $0-24
	MOVQ p+0(FP), AX
	MOVQ n+8(FP), CX
	ADDQ CX, AX
	MOVQ AX, ret+16(FP)
	RET
`)
	code, _, err := Assemble(fn)
	if err != nil {
		t.Fatalf("Assemble: %v", err)
	}
	// From `go tool objdump` of the Go-assembled function:
	//   MOVQ 0x8(SP), AX   488b442408
	//   MOVQ 0x10(SP), CX  488b4c2410
	//   ADDQ CX, AX        4801c8
	//   MOVQ AX, 0x18(SP)  4889442418
	//   RET                c3
	want := []byte{
		0x48, 0x8b, 0x44, 0x24, 0x08,
		0x48, 0x8b, 0x4c, 0x24, 0x10,
		0x48, 0x01, 0xc8,
		0x48, 0x89, 0x44, 0x24, 0x18,
		0xc3,
	}
	if hexBytes(code) != hexBytes(want) {
		t.Errorf("FP translation mismatch:\n got:  %s\n want: %s", hexBytes(code), hexBytes(want))
	}
}

// TestAssembleFrame verifies a function with a non-zero frame: the Go-style
// prologue/epilogue and the x+N(FP) → (N+frame+16)(SP) translation, against
// the bytes the Go assembler produces.
func TestAssembleFrame(t *testing.T) {
	fn := firstText(t, `
#include "textflag.h"
TEXT ·withframe(SB), NOSPLIT, $16-16
	MOVQ a+0(FP), AX
	MOVQ b+8(FP), CX
	ADDQ CX, AX
	MOVQ AX, ret+16(FP)
	RET
`)
	code, _, err := Assemble(fn)
	if err != nil {
		t.Fatalf("Assemble: %v", err)
	}
	// From `go tool objdump`:
	//   PUSHQ BP            55
	//   MOVQ SP, BP         4889e5
	//   SUBQ $0x10, SP      4883ec10
	//   MOVQ 0x20(SP), AX   488b442420   (0 + 16 + 16)
	//   MOVQ 0x28(SP), CX   488b4c2428   (8 + 16 + 16)
	//   ADDQ CX, AX         4801c8
	//   MOVQ AX, 0x30(SP)   4889442430   (16 + 16 + 16)
	//   ADDQ $0x10, SP      4883c410
	//   POPQ BP             5d
	//   RET                 c3
	want := []byte{
		0x55, 0x48, 0x89, 0xe5, 0x48, 0x83, 0xec, 0x10,
		0x48, 0x8b, 0x44, 0x24, 0x20,
		0x48, 0x8b, 0x4c, 0x24, 0x28,
		0x48, 0x01, 0xc8,
		0x48, 0x89, 0x44, 0x24, 0x30,
		0x48, 0x83, 0xc4, 0x10, 0x5d, 0xc3,
	}
	if hexBytes(code) != hexBytes(want) {
		t.Errorf("frame translation mismatch:\n got:  %s\n want: %s", hexBytes(code), hexBytes(want))
	}
}

// TestAssembleVexKernel assembles the horizontal-sum reduction the go-flac
// kernels end with; exercising the VEX moves, shuffle and extract forms
// through the full parser → encoder path; and checks the output is
// byte-identical to the Go assembler's.
func TestAssembleVexKernel(t *testing.T) {
	fn := firstText(t, `
#include "textflag.h"
TEXT ·hsum(SB), NOSPLIT, $0
	VPADDQ       Y8, Y9, Y8
	VEXTRACTI128 $1, Y8, X9
	VPADDQ       X9, X8, X8
	VPSHUFD      $0xEE, X8, X9
	VPADDQ       X9, X8, X8
	VMOVQ        X8, AX
	VZEROUPPER
	RET
`)
	code, _, err := Assemble(fn)
	if err != nil {
		t.Fatalf("Assemble: %v", err)
	}
	// From the Go-assembled function:
	//   VPADDQ Y8, Y9, Y8      c44135d4c0
	//   VEXTRACTI128 $1, Y8, X9  c4437d39c101
	//   VPADDQ X9, X8, X8      c44139d4c1
	//   VPSHUFD $0xEE, X8, X9  c4417970c8ee
	//   VPADDQ X9, X8, X8      c44139d4c1
	//   VMOVQ X8, AX           c461f97ec0
	//   VZEROUPPER             c5f877
	//   RET                    c3
	want := []byte{
		0xc4, 0x41, 0x35, 0xd4, 0xc0,
		0xc4, 0x43, 0x7d, 0x39, 0xc1, 0x01,
		0xc4, 0x41, 0x39, 0xd4, 0xc1,
		0xc4, 0x41, 0x79, 0x70, 0xc8, 0xee,
		0xc4, 0x41, 0x39, 0xd4, 0xc1,
		0xc4, 0x61, 0xf9, 0x7e, 0xc0,
		0xc5, 0xf8, 0x77,
		0xc3,
	}
	if hexBytes(code) != hexBytes(want) {
		t.Errorf("VEX kernel mismatch:\n got:  %s\n want: %s", hexBytes(code), hexBytes(want))
	}
}

// TestAssembleShortJumps checks that a tight loop settles on the short (rel8)
// jump forms, byte for byte with the Go assembler.
func TestAssembleShortJumps(t *testing.T) {
	fn := firstText(t, `
#include "textflag.h"
TEXT ·loop(SB), NOSPLIT, $0
	XORQ AX, AX
l1:
	ADDQ $1, AX
	CMPQ AX, $10
	JLT l1
	RET
`)
	code, _, err := Assemble(fn)
	if err != nil {
		t.Fatalf("Assemble: %v", err)
	}
	// From the Go-assembled function:
	//   XORQ AX, AX       4831c0
	//   ADDQ $1, AX       4883c001
	//   CMPQ AX, $10      4883f80a
	//   JLT l1            7cf6      (short, rel8)
	//   RET               c3
	want := []byte{
		0x48, 0x31, 0xc0,
		0x48, 0x83, 0xc0, 0x01,
		0x48, 0x83, 0xf8, 0x0a,
		0x7c, 0xf6,
		0xc3,
	}
	if hexBytes(code) != hexBytes(want) {
		t.Errorf("short-jump mismatch:\n got:  %s\n want: %s", hexBytes(code), hexBytes(want))
	}
}

// TestAssembleJumpFolding checks jump-to-jump folding: a conditional jump to a
// label that only holds an unconditional jump is redirected to the ultimate
// target, exactly as the Go toolchain does before it encodes branches.
func TestAssembleJumpFolding(t *testing.T) {
	fn := firstText(t, `
#include "textflag.h"
TEXT ·fold(SB), NOSPLIT, $0
	XORQ AX, AX
	JGE done
	INCQ AX
done:
	JMP end
end:
	RET
`)
	code, _, err := Assemble(fn)
	if err != nil {
		t.Fatalf("Assemble: %v", err)
	}
	// From the Go-assembled function: the JGE skips past the done: trampoline
	// straight to end:
	//   XORQ AX, AX       4831c0
	//   JGE end           7d05      (folded past done)
	//   INCQ AX           48ffc0
	//   JMP end           eb00
	//   RET               c3
	want := []byte{
		0x48, 0x31, 0xc0,
		0x7d, 0x05,
		0x48, 0xff, 0xc0,
		0xeb, 0x00,
		0xc3,
	}
	if hexBytes(code) != hexBytes(want) {
		t.Errorf("jump-folding mismatch:\n got:  %s\n want: %s", hexBytes(code), hexBytes(want))
	}
}

func TestAssemblePrefetch(t *testing.T) {
	fn := firstText(t, `
#include "textflag.h"
TEXT ·pf(SB), NOSPLIT, $0
	PREFETCHNTA (AX)
	PREFETCHT0 (BX)
	PREFETCHT1 8(CX)
	PREFETCHT2 -1(AX)(R12*1)
	RET
`)
	code, _, err := Assemble(fn)
	if err != nil {
		t.Fatalf("Assemble: %v", err)
	}
	got := strings.Join(disasm(t, code), "\n")
	want := strings.Join([]string{
		"prefetchnta zmmword ptr [rax]",
		"prefetcht0 zmmword ptr [rbx]",
		"prefetcht1 zmmword ptr [rcx+0x8]",
		"prefetcht2 zmmword ptr [rax+r12-0x1]",
		"ret",
	}, "\n")
	if got != want {
		t.Errorf("prefetch disassembly mismatch:\n got:\n%s\n want:\n%s", got, want)
	}
	// Byte-level expectations: 0F 18 with the variant in the reg field.
	if hex := hexBytes(code[:3]); hex != "0f 18 00" {
		t.Errorf("PREFETCHNTA bytes: got %s, want 0f 18 00", hex)
	}
	if hex := hexBytes(code[3:6]); hex != "0f 18 0b" {
		t.Errorf("PREFETCHT0 bytes: got %s, want 0f 18 0b", hex)
	}
}

// TestSubSPEncodings pins the prologue SUB against the bytes go tool asm
// emits for SUBQ $size, SP: imm8 for -128..127, the imm32 form for anything
// larger.  The intermediate 129..255 range used to encode an ADD with a
// truncated immediate, moving SP the wrong way.
func TestSubSPEncodings(t *testing.T) {
	for _, tt := range []struct {
		size int
		want []byte
	}{
		{8, []byte{0x48, 0x83, 0xEC, 0x08}},
		{127, []byte{0x48, 0x83, 0xEC, 0x7F}},
		{128, []byte{0x48, 0x81, 0xEC, 0x80, 0x00, 0x00, 0x00}},
		{200, []byte{0x48, 0x81, 0xEC, 0xC8, 0x00, 0x00, 0x00}},
		{255, []byte{0x48, 0x81, 0xEC, 0xFF, 0x00, 0x00, 0x00}},
		{4096, []byte{0x48, 0x81, 0xEC, 0x00, 0x10, 0x00, 0x00}},
	} {
		got := subSP(tt.size)
		if !bytes.Equal(got, tt.want) {
			t.Errorf("subSP(%d) = %x, want %x", tt.size, got, tt.want)
		}
	}
}