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

package asm

import (
	"strings"
	"testing"

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

// assembleArm64Source parses src and assembles it for arm64, returning the
// first function's words little-endian.
func assembleArm64Source(t *testing.T, src string) []uint32 {
	t.Helper()
	f, errs := parser.Parse("test_arm64.s", src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("AssembleFileARM64: %v", err)
	}
	if len(img.Funcs) != 1 {
		t.Fatalf("got %d funcs, want 1", len(img.Funcs))
	}
	return wordsOf(img.Code[img.Funcs[0].Offset : img.Funcs[0].Offset+img.Funcs[0].Size])
}

func TestArm64FrameAddrEncoding(t *testing.T) {
	src := `#include "textflag.h"

TEXT ·fr(SB), NOSPLIT, $432-24
	MOVD	$argframe+0(FP), R3
	MOVD	$big+4096(FP), R4
	MOVD	$ret-8(FP), R2
	MOVD	$x-64(FP), R5
	MOVD	RSP, R19
	MOVD	R20, RSP
	RET
`
	ws := assembleArm64Source(t, src)
	// Prologue (4: large frame) then the body at words 4..9, the toolchain's
	// own encodings for the same statements:
	//	ADD $456, RSP, R3            (456 = 448 + 8 + 0)
	//	ADD $(1<<12), RSP, R4        (4552, the hi<<12 half)
	//	ADD $456, R4, R4             (then the lo half)
	//	ADD $448, RSP, R2            (448 - 8 + 8)
	//	ADD $392, RSP, R5            (448 - 64 + 8)
	//	ADD $0, RSP, R19             (the SP register move)
	//	ADD $0, R20, RSP
	want := []uint32{
		0xd10703f4, 0xa93ffa9d, 0x9100029f, 0xd10023fd,
		0x910723e3, 0x914007e4, 0x91072084, 0x910703e2,
		0x910623e5, 0x910003f3, 0x9100029f,
	}
	if len(ws) < len(want) {
		t.Fatalf("got %d words, want at least %d", len(ws), len(want))
	}
	for i, w := range want {
		if ws[i] != w {
			t.Errorf("word %d: got %08x, want %08x", i, ws[i], w)
		}
	}
}

func TestArm64FrameAddrSizes(t *testing.T) {
	// One imm12 word inside the addcon band, two inside the 24-bit band.
	tests := []struct {
		v    int64
		verb int
	}{
		{456, 1},
		{0xFFF, 1},
		{0x1000, 1}, // the shifted imm12 form
		{0x1005, 2}, // the hi<<12 plus lo pair
		{0xFFFFFF, 2},
	}
	for _, tt := range tests {
		got := len(arm64FrameAddrWords(tt.v, 3))
		if got != tt.verb {
			t.Errorf("arm64FrameAddrWords(%d) took %d words, want %d", tt.v, got, tt.verb)
		}
	}
	if ws := arm64FrameAddrWords(0x1000, 3); len(ws) != 1 || ws[0] != 0x914007e3 {
		t.Errorf("arm64FrameAddrWords(0x1000) = %08x, want the shifted ADD 914007e3", ws[0])
	}
	if !arm64IsAddcon(0xFFF) || arm64IsAddcon(0x1001) || arm64IsAddcon(-1) {
		t.Error("arm64IsAddcon misclassifies the band edges")
	}
	fp := &ast.Symbol{Name: "x", Pseudo: "FP", Offset: 8}
	if v := arm64FrameAddrValue(fp, arm64FrameInfo{autosize: 448}); v != 464 {
		t.Errorf("FP address: got %d, want 464", v)
	}
	sp := &ast.Symbol{Name: "x", Pseudo: "SP", Offset: 8}
	if v := arm64FrameAddrValue(sp, arm64FrameInfo{frame: 432}); v != 448 {
		t.Errorf("SP address: got %d, want 448", v)
	}
}

func TestArm64FrameAddrRejects(t *testing.T) {
	tests := []struct {
		name string
		src  string
		want string
	}{
		{
			"width", `TEXT ·f(SB), NOSPLIT, $16-8
	MOVW	$x+0(FP), R7
	RET
`, "illegal combination",
		},
		{
			"pool band", `TEXT ·f(SB), NOSPLIT, $20000000-8
	MOVD	$x+20000000(FP), R7
	RET
`, "literal pool",
		},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			f, errs := parser.Parse("test_arm64.s", tt.src)
			if len(errs) > 0 {
				t.Fatalf("parse: %v", errs)
			}
			_, err := AssembleFileARM64(f)
			if err == nil {
				t.Fatalf("%s: no error, want one naming %q", tt.name, tt.want)
			}
			if !strings.Contains(err.Error(), tt.want) {
				t.Errorf("%s: error %q, want it to name %q", tt.name, err, tt.want)
			}
		})
	}
}

func TestArm64SPMoveEncoding(t *testing.T) {
	src := `TEXT ·f(SB), NOSPLIT, $0-8
	MOVD	RSP, R19
	MOVD	R20, RSP
	MOVD	ZR, R4
	MOVD	RSP, RSP
	RET
`
	ws := assembleArm64Source(t, src)
	// The SP register moves ride ADD $0; the zero move stays ORR.
	want := []uint32{
		0x910003f3, // ADD $0, RSP, R19
		0x9100029f, // ADD $0, R20, RSP
		0xaa1f03e4, // ORR R4, ZR, ZR
		0x910003ff, // ADD $0, RSP, RSP
		0xd65f03c0, // RET
	}
	if len(ws) != len(want) {
		t.Fatalf("got %d words, want %d", len(ws), len(want))
	}
	for i, w := range want {
		if ws[i] != w {
			t.Errorf("word %d: got %08x, want %08x", i, ws[i], w)
		}
	}

	rej := `TEXT ·f(SB), NOSPLIT, $0-8
	MOVW	RSP, R7
	RET
`
	f, errs := parser.Parse("test_arm64.s", rej)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	if _, err := AssembleFileARM64(f); err == nil || !strings.Contains(err.Error(), "illegal combination") {
		t.Errorf("MOVW RSP: error %v, want an illegal-combination refusal", err)
	}
}

func TestArm64AliasLiveness(t *testing.T) {
	// A define inside a dead conditional branch must not become an alias:
	// go_tls.h's `#ifdef GOARCH_arm` block defines LR as R14, which would
	// silently renumber the link register on arm64, where LR is R30.
	src := `#define RARG R5
#ifdef GOARCH_arm
#define LR R14
#endif
TEXT ·f(SB), NOSPLIT, $0-8
	MOVD	LR, R0
	MOVD	RARG, R1
	MOVD	R14, R2
	RET
`
	ws := assembleArm64Source(t, src)
	want := []uint32{
		0xaa1e03e0, // ORR R0, ZR, R30: LR stayed the link register
		0xaa0503e1, // ORR R1, ZR, R5: the live alias applied
		0xaa0e03e2, // ORR R2, ZR, R14: R14 is R14
		0xd65f03c0,
	}
	if len(ws) != len(want) {
		t.Fatalf("got %d words, want %d", len(ws), len(want))
	}
	for i, w := range want {
		if ws[i] != w {
			t.Errorf("word %d: got %08x, want %08x", i, ws[i], w)
		}
	}
}

func TestArm64ADRNoChainChase(t *testing.T) {
	// The toolchain's jump-to-jump collapse rewrites branch targets only:
	// a B to a chain-leading label is redirected, an ADR to the same label
	// resolves to the label itself.
	src := `TEXT ·adrchain(SB), NOSPLIT, $0-0
	B	a
a:
	B	b
b:
	ADR	a, R0
	RET
`
	ws := assembleArm64Source(t, src)
	want := []uint32{
		0x14000002, // B +2: chased through a to b
		0x14000001, // B +1: a's own jump to b
		0x10ffffe0, // ADR a, R0: -4, the label itself, unchased
		0xd65f03c0,
	}
	if len(ws) != len(want) {
		t.Fatalf("got %d words, want %d", len(ws), len(want))
	}
	for i, w := range want {
		if ws[i] != w {
			t.Errorf("word %d: got %08x, want %08x", i, ws[i], w)
		}
	}
}