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

package asm

import (
	"bytes"
	"encoding/binary"
	"os"
	"path/filepath"
	"strings"
	"testing"

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

// The mid-function literal pool flush, pinned against the toolchain.  A
// kernel of distinct pooled stores walks the conservative displacement bound
// (asm7.go's maxPCDisp) inside one body, so the pool must drain exactly
// where cmd/internal/obj/arm64's checkpool drains it: the bytes, the flush
// points and the reach of every load literal are the toolchain's own.

// pooledKernel renders a kernel of n distinct pooled stores: each offset
// sits a step of 8 past the aligned split band (every fourth one aligned,
// but 0x2000000+8i+4 never is), so every statement pools its own four-byte
// word and expands to eight instruction bytes, the densest walk towards the
// distance bound an arm64 body can make.  head carries the TEXT line for the
// splitting variant, tail the closing statements.
func pooledKernel(head string, n int, tail string) string {
	var b strings.Builder
	b.WriteString("#include \"textflag.h\"\n")
	if head != "" {
		b.WriteString(head)
		b.WriteString("\n")
	}
	b.WriteString("TEXT \u00b7poolmid(SB), NOSPLIT, $0-0\n")
	for i := 0; i < n; i++ {
		b.WriteString("\tMOVD\tR1, ")
		b.WriteString(decimal(0x2000000 + 8*i + 4))
		b.WriteString("(R2)\n")
	}
	b.WriteString(tail)
	return b.String()
}

// decimal formats v in decimal.
func decimal(v int) string {
	if v == 0 {
		return "0"
	}
	var buf [20]byte
	i := len(buf)
	for v > 0 {
		i--
		buf[i] = byte('0' + v%10)
		v /= 10
	}
	return string(buf[i:])
}

// assemblePooled parses and assembles a generated kernel, returning the
// image and the single function's code words.
func assemblePooled(t *testing.T, head string, n int, tail string) (*Image, []uint32) {
	t.Helper()
	src := pooledKernel(head, n, tail)
	dir := t.TempDir()
	path := filepath.Join(dir, "poolmid_arm64.s")
	if err := os.WriteFile(path, []byte(src), 0o644); err != nil {
		t.Fatal(err)
	}
	f, errs := parser.Parse(path, 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("functions = %d, want 1", len(img.Funcs))
	}
	fn := img.Funcs[0]
	return img, leWords(img.Code[fn.Offset : fn.Offset+fn.Size])
}

// a64BranchWord reports whether w is an unconditional B: op 000101 in bits
// 31..26, the encoding the flush guards ride.
func a64BranchWord(w uint32) bool {
	return w>>26 == 0x05
}

// a64BranchTarget decodes a B word's target byte offset from its own pc.
func a64BranchTarget(w uint32, pc int) int {
	d := int(w & 0x03FFFFFF)
	if d&(1<<25) != 0 {
		d |= ^0x03FFFFFF
	}
	return pc + 4*d
}

// a64LoadLiteral reports whether w is a LDR W/X literal (the pool's loads
// into REGTMP: word forms 0x18 and 0x58 in the top byte) and decodes its
// target byte displacement.
func a64LoadLiteral(w uint32) (int, bool) {
	if w>>24 != 0x18 && w>>24 != 0x58 {
		return 0, false
	}
	d := int(w>>5) & 0x7FFFF
	if d&(1<<18) != 0 {
		d |= ^0x7FFFF
	}
	return d * 4, true
}

// TestArm64PoolFlushStructure assembles the distance-bound kernel and pins
// the flush structure on gasm's own image: exactly one mid-body branch over
// a drained segment, no code words inside it, every load literal inside the
// conservative displacement bound, and the segment's byte range free of line
// rows: the pool words carry the flushing statement's source line, so the
// pc-line tables see no delta across them.
func TestArm64PoolFlushStructure(t *testing.T) {
	const n = 44000 // one flush: the bound arrives at about 43690 statements
	img, words := assemblePooled(t, "", n, "\tRET\n")
	fn := img.Funcs[0]

	var branches []int
	for i, w := range words {
		if a64BranchWord(w) {
			branches = append(branches, i*4)
		}
	}
	if len(branches) != 1 {
		t.Fatalf("branch words = %d, want exactly the one flush guard", len(branches))
	}
	branchPC := branches[0]
	target := a64BranchTarget(words[branchPC/4], branchPC)
	segStart, segEnd := branchPC+4, target
	if segEnd <= segStart || segEnd%4 != 0 {
		t.Fatalf("flush branch target %d leaves no legal word range after %d", target, branchPC)
	}
	if segEnd+4 > len(words)*4 {
		t.Fatalf("flush branch target %d runs past the image (%d words)", target, len(words))
	}
	// The drained segment holds pool words only: four-byte constants, no
	// branch opcodes, no literal loads.
	for off := segStart; off < segEnd; off += 4 {
		w := words[off/4]
		if a64BranchWord(w) {
			t.Fatalf("word at %d inside the drained segment is a branch: %08x", off, w)
		}
		if _, lit := a64LoadLiteral(w); lit {
			t.Fatalf("word at %d inside the drained segment is a literal load: %08x", off, w)
		}
	}

	// Every load literal resolves inside the conservative bound and inside
	// the image: the property the flush exists to keep.
	for i, w := range words {
		d, ok := a64LoadLiteral(w)
		if !ok {
			continue
		}
		pc := i * 4
		if pc+d < 0 || pc+d >= len(words)*4 {
			t.Fatalf("literal load at %d targets %d, outside the image", pc, pc+d)
		}
		if d >= a64MaxPCDisp || d <= -a64MaxPCDisp {
			t.Fatalf("literal load at %d displaces %d, outside the conservative bound", pc, d)
		}
	}

	// The drained words carry no line rows of their own: the toolchain gives
	// them the flushing statement's Pos so the pc-line tables see no delta.
	for _, l := range fn.Lines {
		if l.Offset >= segStart && l.Offset < segEnd {
			t.Fatalf("line row at %d sits inside the drained segment [%d, %d)", l.Offset, segStart, segEnd)
		}
	}
}

// TestArm64PoolFlushDifferential assembles the kernel family with gasm and
// with the installed toolchain and holds the bytes equal, the toolchain's
// tail alignment padding excepted: one flush, a fall-through end behind the
// UNDEF guard, two flushes, and a splitting function whose guard prefix and
// morestack block sit behind a shifted pool.
func TestArm64PoolFlushDifferential(t *testing.T) {
	if testing.Short() {
		t.Skip("live go tool asm oracle: skipped in -short mode")
	}
	for _, k := range []struct {
		name string
		head string
		n    int
		tail string
	}{
		{"one flush, RET end", "", 44000, "\tRET\n"},
		{"one flush, UNDEF end", "", 44000, ""},
		{"two flushes", "", 90000, "\tRET\n"},
		{"split function", "TEXT \u00b7poolmid(SB), $8-0", 44000, "\tRET\n"},
	} {
		t.Run(k.name, func(t *testing.T) {
			if k.head != "" {
				// The split variant spells its own TEXT: the generator's
				// NOSPLIT line must give way to it.
				runPoolFlushCase(t, pooledKernelFor(k.head, k.n, k.tail))
				return
			}
			runPoolFlushCase(t, pooledKernel("", k.n, k.tail))
		})
	}
}

// pooledKernelFor renders the kernel with an explicit TEXT line, the split
// variant's shape: no NOSPLIT, so the frame forces the stack-split guard.
func pooledKernelFor(text string, n int, tail string) string {
	var b strings.Builder
	b.WriteString("#include \"textflag.h\"\n")
	b.WriteString(text)
	b.WriteString("\n")
	for i := 0; i < n; i++ {
		b.WriteString("\tMOVD\tR1, ")
		b.WriteString(decimal(0x2000000 + 8*i + 4))
		b.WriteString("(R2)\n")
	}
	b.WriteString(tail)
	return b.String()
}

// runPoolFlushCase assembles one generated kernel both ways and compares
// the function's bytes, relocation sites masked, the toolchain's trailing
// alignment zeros excepted.
func runPoolFlushCase(t *testing.T, src string) {
	t.Helper()
	dir := t.TempDir()
	path := filepath.Join(dir, "poolmid_arm64.s")
	if err := os.WriteFile(path, []byte(src), 0o644); err != nil {
		t.Fatal(err)
	}
	f, errs := parser.Parse(path, src)
	if len(errs) > 0 {
		t.Fatalf("parse: %v", errs)
	}
	img, err := AssembleFileARM64(f)
	if err != nil {
		t.Fatalf("AssembleFileARM64: %v", err)
	}
	gt := oracleFuncCode(t, toolAsmObject(t, path, "arm64"))
	byLocal := make(map[string][]byte, len(gt))
	for name, code := range gt {
		if _, after, ok := strings.Cut(name, "."); ok {
			name = after
		}
		byLocal[name] = code
	}
	for _, fn := range img.Funcs {
		gasmCode := maskCode(append([]byte(nil), img.Code[fn.Offset:fn.Offset+fn.Size]...), fn.Relocs)
		goCode, ok := byLocal[fn.Name]
		if !ok {
			t.Fatalf("%s: not in the oracle output (%d functions)", fn.Name, len(gt))
		}
		goCode = maskCode(append([]byte(nil), goCode...), fn.Relocs)
		cmpLen := min(len(goCode), len(gasmCode))
		if !bytes.Equal(gasmCode[:cmpLen], goCode[:cmpLen]) {
			for w := 0; w < cmpLen/4; w++ {
				g := binary.LittleEndian.Uint32(gasmCode[w*4:])
				o := binary.LittleEndian.Uint32(goCode[w*4:])
				if g != o {
					t.Fatalf("%s: word %d (offset %d) differs: gasm %08x go %08x", fn.Name, w, w*4, g, o)
				}
			}
			t.Fatalf("%s: prefixes equal but lengths differ (gasm %d, oracle %d)", fn.Name, len(gasmCode), len(goCode))
		}
		for _, b := range goCode[len(gasmCode):] {
			if b != 0 {
				t.Fatalf("%s: non-zero trailing bytes in the oracle output", fn.Name)
			}
		}
	}
}