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

package disasm

import (
	"strings"
	"testing"

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

func TestDecodeKnownBytes(t *testing.T) {
	for _, tt := range []struct {
		a    arch.Arch
		code []byte
		text string
		want int
	}{
		{arch.AMD64, []byte{0x55}, "push rbp", 1},
		{arch.AMD64, []byte{0x48, 0x89, 0xE5}, "mov rbp, rsp", 3},
		{arch.ARM64, []byte{0xc0, 0x03, 0x5f, 0xd6}, "RET", 4},
		{arch.RISCV, []byte{0x67, 0x80, 0x00, 0x00}, "RET", 4},
		{arch.LOONG64, []byte{0x20, 0x00, 0x00, 0x4c}, "RET", 4},
	} {
		ins, err := Decode(tt.a, tt.code, 0)
		if err != nil {
			t.Errorf("%s: %v", tt.a, err)
			continue
		}
		if ins.Text != tt.text || ins.Len != tt.want {
			t.Errorf("%s: % x decoded to %q (%d bytes), want %q (%d)",
				tt.a, tt.code, ins.Text, ins.Len, tt.text, tt.want)
		}
	}
}

func TestDecodeUndecodable(t *testing.T) {
	// Zero words do not encode a usable instruction on arm64 and loong64; the
	// placeholder keeps a listing going.  RISC-V is the exception: an all-zero
	// word is the defined UNIMP instruction.
	for _, a := range []arch.Arch{arch.ARM64, arch.LOONG64} {
		ins, err := Decode(a, []byte{0, 0, 0, 0}, 0)
		if err != nil {
			t.Fatalf("%s: %v", a, err)
		}
		if ins.Text != "???" {
			t.Errorf("%s: text = %q, want ???", a, ins.Text)
		}
	}
	// The compressed quadrant claims the zero halfword first, so the zero
	// word decodes as the 2-byte compressed UNIMP.
	if ins, err := Decode(arch.RISCV, []byte{0, 0, 0, 0}, 0); err != nil || ins.Text != "UNIMP" || ins.Len != 2 {
		t.Errorf("riscv zero word: %q len %d err %v, want UNIMP with 2 bytes", ins.Text, ins.Len, err)
	}
	if _, err := Decode(arch.ARM64, []byte{0, 0}, 0); err == nil {
		t.Error("short input: expected an error")
	}
	if _, err := Decode(arch.AMD64, nil, 0); err == nil {
		t.Error("empty input: expected an error")
	}
}

// TestBlockRoundTrip assembles a small kernel with the gasm encoder for every
// architecture and disassembles it back: the listing must cover the whole
// function and end in RET.
func TestBlockRoundTrip(t *testing.T) {
	for _, tt := range []struct {
		a    arch.Arch
		name string
	}{
		{arch.AMD64, "k_amd64.s"},
		{arch.ARM64, "k_arm64.s"},
		{arch.RISCV, "k_riscv64.s"},
		{arch.LOONG64, "k_loong64.s"},
	} {
		src := "TEXT \u00b7k(SB), NOSPLIT, $0\n\tMOVQ AX, CX\n\tRET\n"
		if tt.a != arch.AMD64 {
			src = "TEXT \u00b7k(SB), NOSPLIT, $0\n\tRET\n"
		}
		f, errs := parser.Parse(tt.name, src)
		if len(errs) > 0 {
			t.Fatalf("%s: parse: %v", tt.a, errs)
		}
		img, err := assemble(t, tt.a, f)
		if err != nil {
			t.Fatalf("%s: assemble: %v", tt.a, err)
		}
		fn := img.Funcs[0]
		code := img.Code[fn.Offset : fn.Offset+fn.Size]

		ins := Block(tt.a, code, 0, 100)
		if len(ins) == 0 {
			t.Fatalf("%s: empty listing", tt.a)
		}
		consumed := 0
		for _, in := range ins {
			if in.Text == "" || in.Text == "???" {
				t.Errorf("%s: undecoded instruction at %#x: %q", tt.a, in.Addr, in.Text)
			}
			consumed += in.Len
		}
		if consumed != len(code) {
			t.Errorf("%s: listing consumed %d of %d bytes", tt.a, consumed, len(code))
		}
		if last := ins[len(ins)-1]; !strings.Contains(strings.ToLower(last.Text), "ret") {
			t.Errorf("%s: last instruction = %q, want RET", tt.a, last.Text)
		}
	}
}

func TestBlockLimits(t *testing.T) {
	code := []byte{0x55, 0x55, 0x55, 0x55, 0x55}
	if got := Block(arch.AMD64, code, 0, 3); len(got) != 3 {
		t.Errorf("max=3 produced %d instructions, want 3", len(got))
	}
	if got := Block(arch.AMD64, code, 0, 100); len(got) != 5 {
		t.Errorf("code end produced %d instructions, want 5", len(got))
	}
	if got := Block(arch.AMD64, nil, 0, 3); len(got) != 0 {
		t.Errorf("empty code produced %d instructions, want 0", len(got))
	}
}

// assemble assembles the parsed file with the encoder for a.
func assemble(t *testing.T, a arch.Arch, f *ast.File) (*asm.Image, error) {
	t.Helper()
	switch a {
	case arch.ARM64:
		return asm.AssembleFileARM64(f)
	case arch.RISCV:
		return asm.AssembleFileRISCV(f)
	case arch.LOONG64:
		return asm.AssembleFileLOONG64(f)
	default:
		return asm.AssembleFile(f)
	}
}