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

package asm

import (
	"encoding/binary"
	"strings"
	"testing"

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

// assembleArm64Words parses src, assembles it for arm64 and returns the
// first function's body as little-endian instruction words.  Every statement
// must encode: a failure is the test's.
func assembleArm64Words(t *testing.T, src string) []uint32 {
	t.Helper()
	f, errs := parser.Parse("ext_arm64.s", src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("assemble: %v", err)
	}
	if len(img.Funcs) != 1 {
		t.Fatalf("got %d functions, want 1", len(img.Funcs))
	}
	body := img.Code[img.Funcs[0].Offset:][:img.Funcs[0].Size]
	if len(body)%4 != 0 {
		t.Fatalf("body is %d bytes, not a whole number of instructions", len(body))
	}
	words := make([]uint32, 0, len(body)/4)
	for i := 0; i < len(body); i += 4 {
		words = append(words, binary.LittleEndian.Uint32(body[i:]))
	}
	return words
}

// assembleArm64SrcError parses and assembles src and returns the assembler's
// error text.
func assembleArm64SrcError(t *testing.T, src string) string {
	t.Helper()
	f, errs := parser.Parse("ext_arm64.s", src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	_, err := AssembleFileARM64(f)
	if err == nil {
		t.Fatal("assembled, want an error")
	}
	return err.Error()
}

const arm64ExtProbeHead = "#include \"textflag.h\"\nTEXT ·t(SB), NOSPLIT, $0\n"

// TestArm64AssembleExtensionGolden drives the wired layer through the full
// assembler: text in, instruction word out.  Each want is the encoding the
// ARM Architecture Reference Manual's field layout gives for the statement:
// the fixed class word, the size field from the arrangement, and the
// register and immediate fields in the class's own places.  The arch-level
// golden vectors in arch/arm64_ext_test.go pin the same bytes at the
// metadata layer; these pin the text-to-bytes path.
func TestArm64AssembleExtensionGolden(t *testing.T) {
	tests := []struct {
		stmt string
		want uint32
	}{
		// Unpredicated three-vector: Zn, Zm, Zd, one shared arrangement.
		{"ADD Z1.S, Z2.S, Z0.S", 0x04a20020},
		{"ADD Z0.B, Z1.B, Z2.B", 0x04210002},
		{"SUB Z31.D, Z30.D, Z29.D", 0x04fe07fd},
		{"SQADD Z5.H, Z6.H, Z7.H", 0x046610a7},
		{"UQADD Z8.S, Z9.S, Z10.S", 0x04a9150a},
		{"SQSUB Z5.H, Z6.H, Z7.H", 0x046618a7},
		{"UQSUB Z8.S, Z9.S, Z10.S", 0x04a91d0a},
		{"MUL Z0.B, Z1.B, Z2.B", 0x04216002},
		{"SMULH Z11.D, Z12.D, Z13.D", 0x04ec696d},
		{"UMULH Z0.B, Z1.B, Z2.B", 0x04216c02},
		// Governed destructive, merging: Zm, Pg/M, Zdn; the governing
		// predicate is a 3-bit field, so P0-P7 alone.
		{"ADD Z1.S, P0/M, Z0.S", 0x04800020},
		{"SUBR Z1.S, P7/M, Z0.S", 0x04831c20},
		{"MUL Z3.D, P2/M, Z5.D", 0x04d00865},
		{"SUBR Z0.B, P5/M, Z31.B", 0x0403141f},
		// Immediate classes: imm{, LSL #8}, Zdn.  A bare multiple of 256
		// derives the shift bit, the spelling the layer canonicalises.
		{"ADD $255, Z0.S", 0x25a0dfe0},
		{"ADD $65280, Z0.H", 0x2560ffe0},
		{"ADD $255, LSL #8, Z0.S", 0x25a0ffe0},
		{"ADD $(255<<8), Z0.S", 0x25a0ffe0},
		{"MUL $-128, Z0.B", 0x2530d000},
	}
	for _, tt := range tests {
		words := assembleArm64Words(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
		if len(words) != 2 {
			t.Fatalf("%s: got %d words, want the statement and RET", tt.stmt, len(words))
		}
		if words[0] != tt.want {
			t.Errorf("%s:\n  got  %08x\n  want %08x", tt.stmt, words[0], tt.want)
		}
		if words[1] != 0xd65f03c0 {
			t.Errorf("%s: RET encoded %08x", tt.stmt, words[1])
		}
	}
}

// TestArm64AssembleExtensionRefusals pins the diagnostics a pinned statement
// gets from the layer instead of a scalar path's register complaint.
func TestArm64AssembleExtensionRefusals(t *testing.T) {
	tests := []struct {
		stmt string
		want string
	}{
		{"ADD Z0, Z1.S, Z2.S", "carries no arrangement suffix"},
		{"ADD Z33.S, Z1.S, Z2.S", "outside Z0-Z31"},
		{"ADD Z1.S, P0/Z, Z0.S", "/M"},
		{"ADD Z1.S, P9/M, Z0.S", "outside P0-P7"},
		{"ADD $300, Z0.S", "immediate 300"},
		{"ADD $255<<8, Z0.S", "not an immediate the layer can read"},
		{"ADD Z1.S, P0/M, R0", "wants a scalable vector register"},
		{"ADD Z1.S, P0/B, Z0.S", "not an extended-layer operand"},
		{"MUL $200, Z0.B", "outside the signed 8-bit range"},
		{"ADD Z1.S, LSL #8, Z0.S", "LSL belongs straight after an immediate"},
	}
	for _, tt := range tests {
		got := assembleArm64SrcError(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
		if !strings.Contains(got, tt.want) {
			t.Errorf("%s: error %q does not name %q", tt.stmt, got, tt.want)
		}
	}
}

// TestArm64AssembleExtensionPredicateGolden drives the predicate family
// through the full assembler: corpus spellings in, corpus words out.  The
// want words are the same arm64sveenc.s lines the arch-level golden test
// pins; these prove the text-to-bytes path parses the dot and slash
// qualifiers, the general registers and the counter spelling the statements
// write.
func TestArm64AssembleExtensionPredicateGolden(t *testing.T) {
	tests := []struct {
		stmt string
		want uint32
	}{
		// The logical operations, .B alone, the governing predicate with the
		// zeroing qualifier in its dot spelling.
		{"PAND P4.B, P2.B, P1.Z, P14.B", 0x2504444e},
		{"PANDS P4.B, P2.B, P1.Z, P14.B", 0x2544444e},
		{"PBIC P4.B, P2.B, P1.Z, P14.B", 0x2504445e},
		{"PEOR P4.B, P2.B, P1.Z, P14.B", 0x2504464e},
		{"PNAND P4.B, P2.B, P1.Z, P14.B", 0x2584465e},
		{"PORR P4.B, P2.B, P1.Z, P14.B", 0x2584444e},
		// The slash qualifier spells the same operand the dot spelling does.
		{"PAND P4.B, P2.B, P1/Z, P14.B", 0x2504444e},
		// The select and the breaks.
		{"PSEL P4.B, P2.B, P1, P14.B", 0x2504465e},
		{"PBRKA P5.B, P9.Z, P2.B", 0x251064a2},
		{"PBRKAS P5.B, P9.Z, P4.B", 0x255064a4},
		{"PBRKN P4.B, P2.B, P1.Z, P4.B", 0x25184444},
		{"PBRKPA P4.B, P2.B, P1.Z, P14.B", 0x2504c44e},
		// The permutations carry the arrangement into the size field.
		{"PTRN1 P5.D, P4.D, P2.D", 0x05e55082},
		{"PUZP2 P5.D, P4.D, P2.D", 0x05e54c82},
		{"PZIP1 P5.H, P4.H, P2.H", 0x05654082},
		// The singles and the first-fault group.
		{"PPFALSE P13.B", 0x2518e40d},
		{"PPFIRST P5.B, P9, P5.B", 0x2558c125},
		{"PPNEXT P5.D, P4, P5.D", 0x25d9c485},
		{"PPTEST P14.B, P0", 0x2550c1c0},
		{"PPUNPKHI P14.B, P0.H", 0x053141c0},
		{"PRDFFR P13.B", 0x2519f00d},
		{"PRDFFR P14.Z, P0.B", 0x2518f1c0},
		{"PRDFFRS P14.Z, P0.B", 0x2558f1c0},
		{"PWRFFR P13.B", 0x252891a0},
		{"PREV P14.S, P13.S", 0x05b441cd},
		{"SETFFR", 0x252c9000},
		// The while compares: general registers in, a sized predicate out.
		{"PWHILEGE R2, R10, P10.H", 0x2562114a},
		{"PWHILELT R2, R10, P10.H", 0x2562154a},
		{"PWHILELS R2, R10, P10.B", 0x25221d5a},
		{"PWHILERW R2, R10, P10.H", 0x2562315a},
	}
	for _, tt := range tests {
		words := assembleArm64Words(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
		if len(words) != 2 {
			t.Fatalf("%s: got %d words, want the statement and RET", tt.stmt, len(words))
		}
		if words[0] != tt.want {
			t.Errorf("%s:\n  got  %08x\n  want %08x", tt.stmt, words[0], tt.want)
		}
		if words[1] != 0xd65f03c0 {
			t.Errorf("%s: RET encoded %08x", tt.stmt, words[1])
		}
	}
}

// TestArm64AssembleExtensionPredicateRefusals pins the diagnostics the
// predicate statements get from the layer.
func TestArm64AssembleExtensionPredicateRefusals(t *testing.T) {
	tests := []struct {
		stmt string
		want string
	}{
		{"PAND P4.B, P2.B, P1.M, P14.B", "zeroing qualifier"},
		{"PAND P4.B, P2.B, P9.Z, P14.B", "outside P0-P7"},
		{"PAND P8.B, P2.B, P1.Z, P14.B", "outside P0-P7"},
		{"PAND P4.S, P2.B, P1.Z, P14.B", "want .B"},
		{"PBRKN P4.B, P2.B, P1.Z, P9.B", "same register Pdm"},
		{"PBRKAS P5.B, P9.M, P4.B", "zeroing qualifier"},
		{"PSEL P4.B, P2.B, P1.Z, P14.B", "takes no qualifier"},
		{"PTRN1 P5.D, P4.S, P2.D", "want .D"},
		{"PPFALSE P13.S", "want .B"},
		{"PPTRUE P6.S", "predicate-as-counter register"},
		{"PPTRUE PN6.S", "outside PN8-PN15"},
		{"PWHILELT R2, R31, P10.H", "outside R0-R30"},
		{"PWHILELT R2, R10, P10.Q", "no size encoding"},
		{"SETFFR P0.B", "takes 0 operands"},
	}
	for _, tt := range tests {
		got := assembleArm64SrcError(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
		if !strings.Contains(got, tt.want) {
			t.Errorf("%s: error %q does not name %q", tt.stmt, got, tt.want)
		}
	}
}

// TestArm64AssembleExtensionPredicateLeavesScalarsAlone pins the non-
// invasion promise across the new operand kinds: general-register
// statements whose mnemonics the layer also carries for predicates keep
// their scalar behaviour whenever no vector, predicate or counter operand
// appears.
func TestArm64AssembleExtensionPredicateLeavesScalarsAlone(t *testing.T) {
	tests := []struct {
		stmt string
		want uint32
	}{
		{"ADD R0, R1, R2", 0x8b000022},
		{"SUB R0, R1, R2", 0xcb000022},
	}
	for _, tt := range tests {
		words := assembleArm64Words(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
		if len(words) != 2 {
			t.Fatalf("%s: got %d words, want the statement and RET", tt.stmt, len(words))
		}
		if words[0] != tt.want {
			t.Errorf("%s:\n  got  %08x\n  want %08x", tt.stmt, words[0], tt.want)
		}
	}
}

// TestArm64AssembleExtensionPredicateLabelOffsets proves pass 1 and pass 2
// agree on a function mixing the predicate family with the vector family:
// PWHILELT (4 bytes) and SETFFR (4 bytes) shift the label by exactly the
// words pass 2 lays down.
func TestArm64AssembleExtensionPredicateLabelOffsets(t *testing.T) {
	src := arm64ExtProbeHead + `
	PWHILELT R2, R10, P10.H
loop:
	SETFFR
	B loop
	PPFALSE P13.B
	RET
`
	words := assembleArm64Words(t, src)
	want := []uint32{0x2562154a, 0x252c9000, 0x17ffffff, 0x2518e40d, 0xd65f03c0}
	if len(words) != len(want) {
		t.Fatalf("got %d words, want %d", len(words), len(want))
	}
	for i := range want {
		if words[i] != want[i] {
			t.Errorf("word %d: got %08x, want %08x", i, words[i], want[i])
		}
	}
}

// TestArm64AssembleExtensionLabelOffsets proves pass 1 and pass 2 agree on a
// function that mixes the layer with ordinary statements: the label after an
// SVE instruction lands on the 4 bytes the encoder laid down, and the branch
// back to it encodes the distance in words.
func TestArm64AssembleExtensionLabelOffsets(t *testing.T) {
	src := arm64ExtProbeHead + `
	ADD Z0.S, Z1.S, Z2.S
loop:
	ADD $255, Z0.S
	B loop
	MUL $-128, Z0.B
	RET
`
	words := assembleArm64Words(t, src)
	want := []uint32{0x04a10002, 0x25a0dfe0, 0x17ffffff, 0x2530d000, 0xd65f03c0}
	if len(words) != len(want) {
		t.Fatalf("got %d words, want %d", len(words), len(want))
	}
	for i := range want {
		if words[i] != want[i] {
			t.Errorf("word %d: got %08x, want %08x", i, words[i], want[i])
		}
	}
}

// TestArm64AssembleExtensionLeavesScalarsAlone pins the non-invasion
// promise: statements whose operands the scalar paths already read keep
// their exact encodings, scalar and NEON alike, with the layer wired in.
func TestArm64AssembleExtensionLeavesScalarsAlone(t *testing.T) {
	tests := []struct {
		stmt string
		want uint32
	}{
		{"ADD R0, R1, R2", 0x8b000022},
		{"ADD $255, R0", 0x9103fc00},
		{"SUB R0, R1, R2", 0xcb000022},
	}
	for _, tt := range tests {
		words := assembleArm64Words(t, arm64ExtProbeHead+"\t"+tt.stmt+"\n\tRET\n")
		if len(words) != 2 {
			t.Fatalf("%s: got %d words, want the statement and RET", tt.stmt, len(words))
		}
		if words[0] != tt.want {
			t.Errorf("%s:\n  got  %08x\n  want %08x", tt.stmt, words[0], tt.want)
		}
	}
}