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

//go:build linux && amd64

package debug

import (
	"bytes"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"
	"strings"
	"testing"
	"time"
	"unsafe"

	"sourcedock.dev/petrbalvin/gasm-sdk/asm"
	"sourcedock.dev/petrbalvin/gasm-sdk/verify"
)

// Integration tests beyond the basic entry breakpoint: hardware watchpoints,
// conditional breakpoints, next/finish over a CALL, faulting kernels and the
// xstate vector-register readout.  All drive a real ptrace session, so they
// run on amd64 hosts only.

// writeKernel writes an assembly source to a temporary file with the
// architecture suffix the assembler dispatcher expects.
func writeKernel(t *testing.T, src string) string {
	t.Helper()
	path := filepath.Join(t.TempDir(), "kernel_amd64.s")
	if err := os.WriteFile(path, []byte(src), 0o644); err != nil {
		t.Fatalf("write kernel: %v", err)
	}
	return path
}

// launchKernel launches a session for the kernel source and returns the
// session, its breakpoint manager and the function layout.
func launchKernel(t *testing.T, bin, path, funcName string, args []byte) (*Session, *Breakpoints, asm.FuncLayout) {
	t.Helper()
	k, err := verify.Load(path)
	if err != nil {
		t.Fatalf("Load: %v", err)
	}
	t.Cleanup(k.Close)
	fl, err := k.Func(funcName)
	if err != nil {
		t.Fatalf("Func: %v", err)
	}
	if len(args) < fl.Args {
		padded := make([]byte, fl.Args)
		copy(padded, args)
		args = padded
	}
	sess, err := Launch(bin, path, funcName, args)
	if err != nil {
		t.Fatalf("Launch: %v", err)
	}
	t.Cleanup(sess.Kill)
	bm := NewBreakpoints(sess)
	return sess, bm, fl
}

// runToEntry resumes the freshly launched debuggee until the breakpoint at
// the function entry traps, mirroring the REPL continue loop: the debuggee
// SIGSTOPs twice (launch barrier and entry barrier) before entering the JIT
// call.
func runToEntry(t *testing.T, sess *Session, bm *Breakpoints, entry uint64) {
	t.Helper()
	for range 50 {
		for _, bp := range bm.All() {
			bm.Reinsert(bp.Addr)
		}
		if err := sess.Continue(); err != nil {
			t.Fatalf("Continue: %v", err)
		}
		if sess.Exited() {
			t.Fatal("debuggee exited before the entry breakpoint trapped")
		}
		regs, err := sess.GetRegs()
		if err != nil {
			t.Fatalf("GetRegs: %v", err)
		}
		if bm.HandleTrap(&regs) != nil {
			return
		}
	}
	t.Fatal("no entry breakpoint trap after 50 resumes")
}

// captureStdout runs fn with os.Stdout redirected to a pipe and returns
// what it printed (the REPL writes its reports to stdout).
func captureStdout(t *testing.T, fn func()) string {
	t.Helper()
	r, w, err := os.Pipe()
	if err != nil {
		t.Fatalf("pipe: %v", err)
	}
	old := os.Stdout
	os.Stdout = w
	done := make(chan string, 1)
	go func() {
		b, _ := io.ReadAll(r)
		done <- string(b)
	}()
	defer func() { os.Stdout = old }()
	fn()
	w.Close()
	return <-done
}

// TestWatchpointArmRunHit proves the debug-register offsets: the watchpoint
// must fire on the store, with si_addr naming the watched address.  The
// kernel writes its return value to ret+0(FP), which is the 8-byte word
// right above the stack pointer at entry.
func TestWatchpointArmRunHit(t *testing.T) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bin := buildGasm(t)

	const kernel = `#include "textflag.h"

// func wpret() int64
TEXT ·wpret(SB), NOSPLIT, $0-8
	MOVQ $0x5a5a5a5a5a5a5a5a, AX
	MOVQ AX, ret+0(FP)
	RET
`
	path := writeKernel(t, kernel)
	sess, bm, fl := launchKernel(t, bin, path, "wpret", nil)

	entry := sess.CodeBase() + uint64(fl.Offset)
	if _, err := bm.Set(entry, "entry"); err != nil {
		t.Fatalf("Set: %v", err)
	}
	runToEntry(t, sess, bm, entry)

	regs, err := sess.GetRegs()
	if err != nil {
		t.Fatalf("GetRegs: %v", err)
	}
	watched := regs.RSP + 8 // ret+0(FP): the store target

	slot := sess.FindFreeWatchpointSlot()
	if slot < 0 {
		t.Fatal("no free watchpoint slot")
	}
	if err := sess.SetWatchpoint(slot, watched, WatchWrite, 8); err != nil {
		t.Fatalf("SetWatchpoint: %v (wrong debug-register offsets?)", err)
	}

	if err := sess.Continue(); err != nil {
		t.Fatalf("Continue: %v", err)
	}
	reason, addr := sess.StopInfo()
	if reason != StopWatchpoint {
		t.Fatalf("stop reason = %v, want StopWatchpoint (DR0-DR3/DR7 offsets are wrong)", reason)
	}
	if addr != watched {
		t.Fatalf("watchpoint address = %#x, want %#x", addr, watched)
	}

	// The watched word holds the stored value: x86 data breakpoints are
	// reported with the access complete.
	if word, err := sess.Peek(watched); err != nil || word != 0x5a5a5a5a5a5a5a5a {
		t.Errorf("watched word = %#x (err %v), want 0x5a5a5a5a5a5a5a5a", word, err)
	}
	if err := sess.ClearWatchpoint(slot); err != nil {
		t.Fatalf("ClearWatchpoint: %v", err)
	}
}

// TestConditionalBreakpointFalseThenTrue proves the false-condition path:
// the breakpoint steps over the original instruction, re-arms itself and
// keeps running silently, and the true condition stops exactly once with the
// register in the expected state.
func TestConditionalBreakpointFalseThenTrue(t *testing.T) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bin := buildGasm(t)

	const kernel = `#include "textflag.h"

// func countdown(n int64) int64
TEXT ·countdown(SB), NOSPLIT, $0-16
	MOVQ n+0(FP), CX
loop:
	DECQ CX
	CMPQ CX, $0
	JNE loop
	MOVQ CX, ret+8(FP)
	RET
`
	path := writeKernel(t, kernel)
	sess, bm, fl := launchKernel(t, bin, path, "countdown", []byte{8})

	loopAddr := sess.CodeBase() + uint64(fl.Offset) + uint64(fl.Labels["loop"])
	// The length of the breakpointed instruction, from a disassembly taken
	// before the INT3 is patched in.
	_, insnLen, err := sess.Disassemble(loopAddr)
	if err != nil || insnLen <= 0 {
		t.Fatalf("Disassemble at %#x: len=%d err=%v", loopAddr, insnLen, err)
	}
	cond := &Condition{Reg: "rcx", Op: "==", Value: 1}
	bp, err := bm.SetWithCond(loopAddr, "loop", cond)
	if err != nil {
		t.Fatalf("SetWithCond: %v", err)
	}

	hits := 0
	exited := false
	for range 200 {
		for _, b := range bm.All() {
			bm.Reinsert(b.Addr)
		}
		if err := sess.Continue(); err != nil {
			exited = true
			break // the debuggee finished
		}
		if sess.Exited() {
			exited = true
			break
		}
		if sig := sess.LastSignal(); sig != 0 {
			t.Fatalf("unexpected signal stop %v", sig)
		}
		regs, err := sess.GetRegs()
		if err != nil {
			t.Fatalf("GetRegs: %v", err)
		}
		if hit := bm.HandleTrap(&regs); hit != nil {
			hits++
			if regs.RCX != 1 {
				t.Fatalf("hit with RCX=%d, want 1", regs.RCX)
			}
			// Park after the instruction, as the REPL does.
			if err := sess.Step(); err != nil {
				t.Fatalf("Step: %v", err)
			}
		} else {
			// A false evaluation must leave the debuggee past the whole
			// original instruction: a PC inside it (trapAddr+1 on amd64)
			// means the resume happens mid-instruction.
			fresh, err := sess.GetRegs()
			if err != nil {
				t.Fatalf("GetRegs: %v", err)
			}
			if fresh.RIP > loopAddr && fresh.RIP < loopAddr+uint64(insnLen) {
				t.Fatalf("false evaluation left the PC at %#x, inside the %d-byte instruction at %#x",
					fresh.RIP, insnLen, loopAddr)
			}
		}
	}
	if hits != 1 {
		t.Fatalf("conditional breakpoint hit %d times, want exactly 1 (false evaluations must run through silently)", hits)
	}
	if bp.Hits() != 1 {
		t.Errorf("bp.Hits() = %d, want 1", bp.Hits())
	}
	if !exited || !sess.Exited() {
		t.Fatal("debuggee did not run to completion after the conditional hit")
	}
}

// TestNextAndFinishOverCall proves next and finish evaluate the trap with
// registers fetched after the stop: next lands exactly on the instruction
// after the CALL, and finish stops exactly on the return address.
func TestNextAndFinishOverCall(t *testing.T) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bin := buildGasm(t)

	const kernel = `#include "textflag.h"

// func caller(x int64) int64
// The argument travels in AX: FP argument slots of CALL-bearing functions
// are an assembler concern outside this test's scope.
TEXT ·caller(SB), NOSPLIT, $0-16
	MOVQ $5, AX
	CALL ·bump(SB)
aftercall:
	MOVQ AX, ret+8(FP)
	RET

// func bump(x int64) int64
TEXT ·bump(SB), NOSPLIT, $0-0
	ADDQ $3, AX
	RET
`
	path := writeKernel(t, kernel)

	// next: step the prologue and the constant load (3 instructions), then
	// step over the CALL and check the landing address and RAX.
	sess, bm, fl := launchKernel(t, bin, path, "caller", nil)
	entry := sess.CodeBase() + uint64(fl.Offset)
	if _, err := bm.Set(entry, "entry"); err != nil {
		t.Fatalf("Set: %v", err)
	}
	runToEntry(t, sess, bm, entry)
	afterOff := uint64(fl.Labels["aftercall"])

	out := captureStdout(t, func() {
		REPL(sess, bm, sess.CodeBase(), fl.Offset, fl.Size, fl.Args, nil, nil,
			strings.NewReader("step 3\nnext\nregs\nquit\n"))
	})
	if !strings.Contains(out, fmt.Sprintf("func+%#x", afterOff)) {
		t.Errorf("next did not land on the instruction after the CALL (func+%#x); output:\n%s", afterOff, out)
	}
	if !strings.Contains(out, "RAX = 0x0000000000000008") {
		t.Errorf("callee did not run exactly once under next (want RAX=8); output:\n%s", out)
	}

	// finish: run to the return address read off the stack at entry.
	sess2, bm2, fl2 := launchKernel(t, bin, path, "caller", nil)
	entry2 := sess2.CodeBase() + uint64(fl2.Offset)
	if _, err := bm2.Set(entry2, "entry"); err != nil {
		t.Fatalf("Set: %v", err)
	}
	runToEntry(t, sess2, bm2, entry2)
	regs, err := sess2.GetRegs()
	if err != nil {
		t.Fatalf("GetRegs: %v", err)
	}
	retAddr, err := sess2.Peek(regs.RSP)
	if err != nil {
		t.Fatalf("Peek return address: %v", err)
	}

	out2 := captureStdout(t, func() {
		REPL(sess2, bm2, sess2.CodeBase(), fl2.Offset, fl2.Size, fl2.Args, nil, nil,
			strings.NewReader("step 1\nfinish\nquit\n"))
	})
	want := fmt.Sprintf("finished, now at %#x\n", retAddr)
	if !strings.Contains(out2, want) {
		t.Errorf("finish stopped at the wrong PC; want %q in output:\n%s", want, out2)
	}
}

// TestSignalStopSurfaced proves a faulting kernel surfaces as a reported
// stop instead of an infinite fault loop.  A regression here hangs, so a
// watchdog fails the run rather than letting CI stall.
func TestSignalStopSurfaced(t *testing.T) {
	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bin := buildGasm(t)

	const kernel = `#include "textflag.h"

// func crash() int64
TEXT ·crash(SB), NOSPLIT, $0-8
	XORQ AX, AX
	MOVQ (AX), AX
	MOVQ AX, ret+0(FP)
	RET
`
	path := writeKernel(t, kernel)
	sess, bm, _ := launchKernel(t, bin, path, "crash", nil)

	timer := time.AfterFunc(time.Minute, func() {
		panic("watchdog: the debugger hung on the faulting kernel instead of reporting the signal stop")
	})
	defer timer.Stop()

	out := captureStdout(t, func() {
		REPL(sess, bm, sess.CodeBase(), 0, 0, 0, nil, nil,
			strings.NewReader("continue\nquit\n"))
	})
	if !strings.Contains(out, "stopped on signal") {
		t.Errorf("SIGSEGV did not surface as a reported stop; output:\n%s", out)
	}
	if !sess.Exited() {
		t.Error("debuggee should be killed by quit after the signal stop")
	}
}

// TestGetVectorRegsXState proves the NT_X86_XSTATE readout: the request
// succeeds on a normal process and the XMM halves agree with
// PTRACE_GETFPREGS.
func TestGetVectorRegsXState(t *testing.T) {
	// The FPRegs layout must mirror the kernel's user_fpregs_struct
	// exactly: PTRACE_GETFPREGS fills all 512 bytes, so a short struct
	// overflows the caller's memory.
	if got := unsafe.Sizeof(FPRegs{}); got != 512 {
		t.Fatalf("sizeof(FPRegs) = %d, want 512", got)
	}
	if got := unsafe.Offsetof(FPRegs{}.XMM); got != 160 {
		t.Fatalf("offsetof(FPRegs.XMM) = %d, want 160", got)
	}

	runtime.LockOSThread()
	defer runtime.UnlockOSThread()
	bin := buildGasm(t)

	const kernel = `#include "textflag.h"

// func vprobe() int64
TEXT ·vprobe(SB), NOSPLIT, $0-8
	MOVQ $1, AX
	MOVQ AX, ret+0(FP)
	RET
`
	path := writeKernel(t, kernel)
	sess, bm, fl := launchKernel(t, bin, path, "vprobe", nil)

	entry := sess.CodeBase() + uint64(fl.Offset)
	if _, err := bm.Set(entry, "entry"); err != nil {
		t.Fatalf("Set: %v", err)
	}
	runToEntry(t, sess, bm, entry)

	v, err := sess.GetVectorRegs()
	if err != nil {
		t.Fatalf("GetVectorRegs: %v", err)
	}
	fp, err := sess.GetFPRegs()
	if err != nil {
		t.Fatalf("GetFPRegs: %v", err)
	}
	for i := range 16 {
		if !bytes.Equal(v.YMM[i][:16], fp.XMM[i][:]) {
			t.Errorf("YMM%d low half %x, want the FPRegs XMM half %x", i, v.YMM[i][:16], fp.XMM[i][:])
		}
	}
}