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

package asm

import (
	"bytes"
	"os"
	"strings"
	"testing"

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

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

// TestAssembleFileStaticData checks the whole-image layout — code, padding
// and the data section — and that the RIP-relative displacements of static
// symbol loads resolve to the right bytes.
func TestAssembleFileStaticData(t *testing.T) {
	f, errs := parser.Parse("d_amd64.s", `
#include "textflag.h"
TEXT ·load(SB), NOSPLIT, $0
	VMOVDQU mask<>(SB), X15
	MOVL small<>(SB), AX
	RET
GLOBL mask<>(SB), RODATA, $16
DATA mask<>+0(SB)/4, $0x80020100
DATA mask<>+4(SB)/4, $0x80050403
DATA mask<>+8(SB)/4, $0x80080706
DATA mask<>+12(SB)/4, $0x800B0A09
GLOBL small<>(SB), RODATA, $4
DATA small<>+0(SB)/4, $0x1234
`)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("AssembleFile: %v", err)
	}

	// Code (15 bytes) + 1 pad byte to align the data section to 16:
	//   VMOVDQU mask<>(SB), X15   c5 7a 6f 3d 08 00 00 00   (disp = 16 − 8)
	//   MOVL small<>(SB), AX      8b 05 12 00 00 00         (disp = 32 − 14)
	//   RET                       c3
	// Data: pad, mask (16 bytes), small (4 bytes).
	want := "c57a6f3d080000008b0512000000c300" +
		"000102800304058006070880090a0b80" +
		"34120000"
	if got := strings.ReplaceAll(hexBytes(img.Bytes()), " ", ""); got != want {
		t.Errorf("image bytes:\n got  %s\n want %s", got, want)
	}
	if img.Symbols["mask"] != 16 || img.Symbols["small"] != 32 {
		t.Errorf("symbol offsets = %v, want mask=16 small=32", img.Symbols)
	}
	if len(img.Funcs) != 1 || img.Funcs[0].Name != "load" || img.Funcs[0].Size != 15 {
		t.Errorf("funcs = %+v", img.Funcs)
	}
}

// TestAssembleFileErrors checks the static-symbol error paths.
func TestAssembleFileErrors(t *testing.T) {
	cases := []struct {
		name string
		src  string
		want string // substring of the error
	}{
		{
			"undefined symbol",
			`
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	VMOVDQU nope<>(SB), X0
	RET
`,
			"undefined symbol",
		},
		{
			"DATA without GLOBL",
			`
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	RET
DATA orphan<>+0(SB)/4, $1
`,
			"no matching GLOBL",
		},
		{
			"DATA exceeds size",
			`
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	RET
GLOBL tiny<>(SB), RODATA, $4
DATA tiny<>+0(SB)/8, $1
`,
			"exceeds GLOBL size",
		},
		{
			"DATA bad width",
			`
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	RET
GLOBL odd<>(SB), RODATA, $4
DATA odd<>+0(SB)/3, $1
`,
			"invalid width",
		},
	}
	for _, c := range cases {
		f, errs := parser.Parse("e_amd64.s", c.src)
		if len(errs) > 0 {
			t.Fatalf("%s: parse: %v", c.name, errs)
		}
		if _, err := AssembleFile(f); err == nil || !strings.Contains(err.Error(), c.want) {
			t.Errorf("%s: error %v, want substring %q", c.name, err, c.want)
		}
	}

	// A static-symbol operand is unresolvable in single-function assembly.
	fn := firstText(t, `
#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $0
	MOVQ x<>(SB), AX
	RET
GLOBL x<>(SB), RODATA, $8
DATA x<>+0(SB)/4, $1
`)
	if _, _, err := Assemble(fn); err == nil || !strings.Contains(err.Error(), "file-level assembly") {
		t.Errorf("single-function SB: error %v, want a file-level-assembly error", err)
	}
}

// TestAssembleGoFlacAVX2Kernel assembles the whole production AVX2 kernel —
// all functions plus the file-local mask24 constant — and checks that every
// static-symbol load resolves to the right bytes in the image.  Skipped when
// the sibling repository is not checked out.
func TestAssembleGoFlacAVX2Kernel(t *testing.T) {
	path := "../../go-libraries/go-flac/avx2_amd64.s"
	if _, err := os.Stat(path); err != nil {
		t.Skip("go-libraries repository not present next to gasm-devkit")
	}
	src, err := os.ReadFile(path)
	if err != nil {
		t.Fatal(err)
	}
	f, errs := parser.Parse(path, string(src))
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFile(f)
	if err != nil {
		t.Fatalf("AssembleFile: %v", err)
	}
	if len(img.Funcs) != 17 {
		t.Errorf("functions = %d, want 17", len(img.Funcs))
	}

	// mask24 as the DATA directives define it.
	mask := []byte{
		0x00, 0x01, 0x02, 0x80, 0x03, 0x04, 0x05, 0x80,
		0x06, 0x07, 0x08, 0x80, 0x09, 0x0a, 0x0b, 0x80,
	}
	image := img.Bytes()
	if got := image[img.Symbols["mask24"] : img.Symbols["mask24"]+16]; !bytes.Equal(got, mask) {
		t.Errorf("mask24 contents %x, want %x", got, mask)
	}

	// Every VMOVDQU mask24<>(SB), X15 (c5 7a 6f 3d + rel32, i.e. a VMOVDQU
	// with a RIP-relative r/m) must land on the mask bytes within the image.
	loads := 0
	for _, fn := range img.Funcs {
		code := img.Code[fn.Offset : fn.Offset+fn.Size]
		for pc := 0; pc < len(code); {
			inst, err := x86asm.Decode(code[pc:], 64)
			if err != nil {
				t.Fatalf("%s: decode at +%d: %v", fn.Name, pc, err)
			}
			// mod=00, rm=101 → RIP-relative.
			if inst.Op == x86asm.VMOVDQU && inst.Len == 8 && code[pc+3]&0xC7 == 0x05 {
				rel := int32(uint32(code[pc+4]) | uint32(code[pc+5])<<8 | uint32(code[pc+6])<<16 | uint32(code[pc+7])<<24)
				target := fn.Offset + pc + 8 + int(rel)
				if !bytes.Equal(image[target:target+16], mask) {
					t.Errorf("%s: mask load at +%d lands on %x, want %x", fn.Name, pc, image[target:target+16], mask)
				}
				loads++
			}
			pc += inst.Len
		}
	}
	if loads != 2 {
		t.Errorf("mask loads found = %d, want 2", loads)
	}
}