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

package asm

import (
	"bytes"
	"encoding/binary"
	"testing"

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

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

// assembleLOONG64Helper assembles one TEXT function and returns its bytes.
func assembleLOONG64Helper(t *testing.T, fn *ast.Text) []byte {
	t.Helper()
	code, _, _, _, _, err := assembleLOONG64(fn)
	if err != nil {
		t.Fatalf("assemble: %v", err)
	}
	return code
}

// wantWords checks that code matches the expected little-endian words.
func wantWords(t *testing.T, code []byte, want ...uint32) {
	t.Helper()
	got := make([]uint32, 0, len(code)/4)
	for i := 0; i+4 <= len(code); i += 4 {
		got = append(got, binary.LittleEndian.Uint32(code[i:]))
	}
	if len(got) != len(want) {
		t.Fatalf("word count = %d, want %d\ncode: % x", len(got), len(want), code)
	}
	for i := range want {
		if got[i] != want[i] {
			t.Errorf("word %d = %08x, want %08x", i, got[i], want[i])
		}
	}
}

func TestLOONG64_add(t *testing.T) {
	fn := firstTextLOONG64(t, `#include "textflag.h"
TEXT ·add(SB), NOSPLIT, $0-24
	MOVV a+0(FP), R4
	MOVV b+8(FP), R5
	ADDV R5, R4, R4
	MOVV R4, ret+16(FP)
	RET
`)
	code := assembleLOONG64Helper(t, fn)
	// 5 instructions: two ld.d, add.d, st.d, jirl r0, r1, 0.
	wantWords(t, code,
		0x28C02064, // ld.d r4, 8(r3)
		0x28C04065, // ld.d r5, 16(r3)
		0x00109484, // add.d r4, r4, r5
		0x29C06064, // st.d r4, 24(r3)
		0x4C000020, // jirl r0, r1, 0
	)
}

func TestLOONG64_arithmetic(t *testing.T) {
	fn := firstTextLOONG64(t, `#include "textflag.h"
TEXT ·arith(SB), NOSPLIT, $0
	ADDV R4, R5, R6
	SUBV R7, R8, R9
	MULV R10, R11, R12
	DIVV R13, R14, R15
	AND R16, R17, R18
	OR R18, R19, R20
	XOR R20, R21, R2
	SLLV R2, R23, R24
	SRLV R24, R25, R26
	SRAV R26, R27, R28
	RET
`)
	code := assembleLOONG64Helper(t, fn)
	wantWords(t, code,
		0x001090A6, // add.d r6, r5, r4
		0x00119D09, // sub.d r9, r8, r7
		0x001DA96C, // mul.d r12, r11, r10
		0x002235CF, // div.d r15, r14, r13
		0x0014C232, // and r18, r17, r16
		0x00154A74, // or r20, r19, r18
		0x0015D2A2, // xor r2, r21, r20
		0x00188AF8, // sll.d r24, r23, r2
		0x0019633A, // srl.d r26, r25, r24
		0x0019EB7C, // sra.d r28, r27, r26
		0x4C000020, // jirl r0, r1, 0
	)
}

func TestLOONG64_immediates(t *testing.T) {
	fn := firstTextLOONG64(t, `#include "textflag.h"
TEXT ·imm(SB), NOSPLIT, $0
	ADDV $42, R4, R5
	ADDV $-8, R6
	AND $0xff, R7, R8
	OR $1, R9, R10
	SGT $100, R13, R14
	SLLV $4, R15, R16
	MOVV $0x12345, R17
	MOVV $0, R18
	MOVW $0, R19
	RET
`)
	code := assembleLOONG64Helper(t, fn)
	wantWords(t, code,
		0x02C0A885, // addi.d r5, r4, 42
		0x02FFE0C6, // addi.d r6, r6, -8
		0x0343FCE8, // andi r8, r7, 0xff
		0x0380052A, // ori r10, r9, 1
		0x020191AE, // slti r14, r13, 100
		0x004111F0, // slli.d r16, r15, 4
		0x14000251, // lu12i.w r17, 0x12
		0x038D1631, // ori r17, r17, 0x345
		0x00150012, // or r18, r0, r0
		0x00170013, // sll.w r19, r0, r0
		0x4C000020, // jirl r0, r1, 0
	)
}

func TestLOONG64_loadStore(t *testing.T) {
	fn := firstTextLOONG64(t, `#include "textflag.h"
TEXT ·mem(SB), NOSPLIT, $0
	MOVV (R4), R5
	MOVV R5, (R6)
	MOVW 8(R7), R8
	MOVB R9, -4(R10)
	MOVV (R11)(R12), R13
	MOVV R14, (R15)(R16)
	RET
`)
	code := assembleLOONG64Helper(t, fn)
	wantWords(t, code,
		0x28C00085, // ld.d r5, 0(r4)
		0x29C000C5, // st.d r5, 0(r6)
		0x288020E8, // ld.w r8, 8(r7)
		0x293FF149, // st.b r9, -4(r10)
		0x380C316D, // ldx.d r13, r11, r12
		0x381C41EE, // stx.d r14, r15, r16
		0x4C000020, // jirl r0, r1, 0
	)
}

func TestLOONG64_branches(t *testing.T) {
	fn := firstTextLOONG64(t, `#include "textflag.h"
TEXT ·br(SB), NOSPLIT, $0
	BEQ R4, R5, done
	BNE R6, R7, skip
	BLT R8, R9, done
	BGE R10, R11, done
	BLTU R12, R13, done
	BGEU R14, R15, done
skip:
	JMP done
done:
	RET
`)
	code := assembleLOONG64Helper(t, fn)
	// skip is at 0x18 (6 words), done at 0x1c.
	wantWords(t, code,
		0x58001C85, // beq r5, r4, +7
		0x5C0018C7, // bne r7, r6, +6
		0x60001509, // blt r9, r8, +5
		0x6400114B, // bge r11, r10, +4
		0x68000D8D, // bltu r13, r12, +3
		0x6C0009CF, // bgeu r15, r14, +2
		0x50000400, // b done (+1, chain-folded through skip)
		0x4C000020, // jirl r0, r1, 0
	)
}

func TestLOONG64_frame(t *testing.T) {
	fn := firstTextLOONG64(t, `#include "textflag.h"
TEXT ·f(SB), NOSPLIT, $32-8
	MOVV R4, R5
	MOVV arg+0(FP), R6
	MOVV R7, local-8(SP)
	MOVV local-8(SP), R8
	MOVV R9, ret+0(FP)
	RET
`)
	code := assembleLOONG64Helper(t, fn)
	// autosize = align8(32+8) = 40; prologue stores LR at -40(SP),
	// opens the frame, stores LR again at 0(SP).  The function is a leaf
	// (no calls), so the epilogue skips the LR restore.  FP args are at
	// autosize+8; SP locals at autosize+offset.
	wantWords(t, code,
		0x29FF6061, // st.d r1, -40(r3)
		0x02FF6063, // addi.d r3, r3, -40
		0x29C00061, // st.d r1, 0(r3)
		0x00150085, // or r5, r4, r0
		0x28C0C066, // ld.d r6, 48(r3)     arg+0(FP) → 0+40+8
		0x29C08067, // st.d r7, 32(r3)     local-8(SP) → 40-8
		0x28C08068, // ld.d r8, 32(r3)
		0x29C0C069, // st.d r9, 48(r3)     ret+0(FP) → 0+40+8
		0x02C0A063, // addi.d r3, r3, 40
		0x4C000020, // jirl r0, r1, 0
	)
}

func TestLOONG64_jumpChain(t *testing.T) {
	fn := firstTextLOONG64(t, `#include "textflag.h"
TEXT ·jc(SB), NOSPLIT, $0
	JMP a
a:
	JMP b
b:
	RET
`)
	code := assembleLOONG64Helper(t, fn)
	wantWords(t, code,
		0x50000800, // b +2 (a, chain-folded to b)
		0x50000400, // b +1 (b)
		0x4C000020, // jirl r0, r1, 0
	)
}

func TestLOONG64_dconClasses(t *testing.T) {
	cases := []struct {
		v    int64
		word int // expected word count
	}{
		{0x123456789, 3},        // lu12i.w + ori + lu32i.d
		{-1, 2},                 // addi.d + lu52i.d (the MOV path handles -1 earlier)
		{0x1000000000000, 2},    // addi.w + lu32i.d
		{0x123456789abcdef0, 4}, // full sequence
		{0xFFFFFFFFF, 2},        // lu12i.w + ori
		{0x1234567800000000, 3}, // addi.w + lu32i.d + lu52i.d
	}
	for _, c := range cases {
		if n := len(l64DconMovWords(0, c.v)); n != c.word {
			t.Errorf("0x%x: %d words, want %d", c.v, n, c.word)
		}
	}
}

func TestLOONG64_regNames(t *testing.T) {
	cases := map[string]int{
		"R0": 0, "R31": 31, "F0": 0, "F31": 31, "FCC0": 0, "FCC7": 7,
		"FCSR0": 0, "FCSR3": 3, "ZERO": 0, "RA": 1, "SP": 3, "g": 22, "G": 22,
		"R32": -1, "FCC8": -1, "X0": -1, "R": -1, "TMP": 30, "CTXT": 29,
	}
	for name, want := range cases {
		if got := loong64RegNum(name); got != want {
			t.Errorf("loong64RegNum(%q) = %d, want %d", name, got, want)
		}
	}
}

func TestLOONG64_bytesEqualGroundTruth(t *testing.T) {
	// A spot-check that assembleLOONG64 emits the same bytes the Go
	// toolchain does for a small kernel (the full comparison lives in
	// verify's TestGroundTruthLOONG64).
	src := `#include "textflag.h"
TEXT ·k(SB), NOSPLIT, $0-0
	ADDV R4, R5, R6
	MOVV $0x100000, R7
	BEQ R6, R7, done
	JMP done
done:
	RET
`
	fn := firstTextLOONG64(t, src)
	code := assembleLOONG64Helper(t, fn)
	want := []byte{
		0xa6, 0x90, 0x10, 0x00, // add.d r6, r5, r4
		0x07, 0x20, 0x00, 0x14, // lu12i.w r7, 0x100
		0xc7, 0x08, 0x00, 0x58, // beq r7, r6, +2 (done)
		0x00, 0x04, 0x00, 0x50, // b +1 (done)
		0x20, 0x00, 0x00, 0x4c, // jirl r0, r1, 0
	}
	if !bytes.Equal(code, want) {
		t.Errorf("code = % x\nwant % x", code, want)
	}
}