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

//go:build linux

package debug

import (
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"syscall"
	"time"
)

// Session is a ptrace debugging session controlling one debuggee process.
type Session struct {
	pid      int
	cmd      *exec.Cmd
	stopped  bool
	exited   bool
	codeBase uint64   // base address of the JIT code in the debuggee
	tmpDir   string   // scratch directory of the session, removed on Kill
	wpSlots  [16]bool // hardware watchpoint slots in use (DR0-DR3, arm64 DBGWVR0-15)
	// lastSignal holds the signal of the most recent stop when that stop
	// was a genuine signal-delivery-stop the caller must see (a fault such
	// as SIGSEGV, SIGBUS, SIGFPE or SIGILL); 0 for breakpoint traps,
	// single-steps, SIGSTOP and suppressed runtime signals.
	lastSignal syscall.Signal
}

// Launch starts the debuggee subprocess (gasm debug --target ...) and
// attaches to it via ptrace.
func Launch(gasmBin, asmPath, funcName string, args []byte) (*Session, error) {
	sess, _, err := LaunchWithBuffers(gasmBin, asmPath, funcName, args, "")
	return sess, err
}

// LaunchWithBuffers is like Launch but also allocates buffers in the debuggee.
//
// It pins the calling goroutine to its OS thread and leaves it pinned: the
// debuggee's PTRACE_TRACEME binds the tracer relation to the forking thread,
// and every ptrace request on the session must come from that same thread.
// All Session methods must therefore be called from the goroutine that
// launched the session (the REPL and coverage loops do exactly that).
func LaunchWithBuffers(gasmBin, asmPath, funcName string, args []byte, bufSpec string) (*Session, []uint64, error) {
	runtime.LockOSThread() // ptrace requests must stay on the forking thread
	self, err := os.Executable()
	if err != nil {
		return nil, nil, fmt.Errorf("debug: cannot find gasm binary: %w", err)
	}
	if gasmBin != "" {
		self = gasmBin
	}

	tmpDir, err := os.MkdirTemp("", "gasm-debug-*")
	if err != nil {
		return nil, nil, fmt.Errorf("debug: tempdir: %w", err)
	}
	argsFile := filepath.Join(tmpDir, "args.bin")
	if err := os.WriteFile(argsFile, args, 0o644); err != nil {
		os.RemoveAll(tmpDir)
		return nil, nil, fmt.Errorf("debug: write args: %w", err)
	}

	if bufSpec != "" {
		if err := os.WriteFile(filepath.Join(tmpDir, "bufspec"), []byte(bufSpec), 0o644); err != nil {
			os.RemoveAll(tmpDir)
			return nil, nil, fmt.Errorf("debug: write bufspec: %w", err)
		}
	}

	cmd := exec.Command(self, "debug", "--func", funcName, "--args", argsFile, asmPath)
	cmd.Env = append(os.Environ(), "GASM_DEBUG_TARGET=1", "GASM_DEBUG_TMP="+tmpDir)
	cmd.Stdout = nil
	cmd.Stderr = os.Stderr
	cmd.SysProcAttr = &syscall.SysProcAttr{}

	if err := cmd.Start(); err != nil {
		os.RemoveAll(tmpDir)
		return nil, nil, fmt.Errorf("debug: start debuggee: %w", err)
	}

	s := &Session{pid: cmd.Process.Pid, cmd: cmd, tmpDir: tmpDir}

	readyFile := filepath.Join(tmpDir, "ready")
	for range 500 {
		if _, err := os.Stat(readyFile); err == nil {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}

	// The debuggee parks itself with SIGSTOP once the JIT code is mapped.
	// A Go tracee also reports SIGURG preemption as signal-delivery-stops,
	// so the wait loops until a stop the debugger cares about instead of
	// assuming the first event is the SIGSTOP.
	if _, err := s.waitStopped(); err != nil {
		cmd.Process.Kill()
		os.RemoveAll(tmpDir)
		return nil, nil, fmt.Errorf("debug: wait for debuggee: %w", err)
	}
	s.stopped = true

	// The debuggee reports its JIT mapping in the codebase file; that is the
	// exact region the kernel was written to.  Scanning /proc/pid/maps for
	// any RWX region is only the fallback.
	if data, err := os.ReadFile(filepath.Join(tmpDir, "codebase")); err == nil {
		fmt.Sscanf(string(data), "%d", &s.codeBase)
	}
	if s.codeBase == 0 {
		s.codeBase = findRWXMapping(s.pid)
	}

	var bufAddrs []uint64
	if bufSpec != "" {
		addrFile := filepath.Join(tmpDir, "bufaddrs")
		if data, err := os.ReadFile(addrFile); err == nil {
			for line := range strings.SplitSeq(strings.TrimSpace(string(data)), "\n") {
				var addr uint64
				if _, err := fmt.Sscanf(line, "%d", &addr); err == nil {
					bufAddrs = append(bufAddrs, addr)
				}
			}
		}
	}

	return s, bufAddrs, nil
}

// waitStopped consumes ptrace-stop events until one the debugger cares
// about arrives: SIGTRAP (a breakpoint or a completed single-step), the
// debuggee's own SIGSTOP, or a genuine signal-delivery-stop.  A Go tracee's
// runtime raises SIGURG for asynchronous preemption, and every signal on a
// traced thread surfaces as a signal-delivery-stop, so SIGURG is suppressed
// and the tracee resumed without it.  Every other signal (SIGSEGV, SIGBUS,
// SIGFPE, SIGILL, ...) is returned to the caller: resuming with signal 0
// would restart the faulting instruction and fault forever, so a faulting
// kernel must surface as a stop the caller reports.  Runtime noise is also
// why a single wait can return in the middle of runtime code and a resume
// can then fail: the event stream must be drained by the tracer.
func (s *Session) waitStopped() (syscall.Signal, error) {
	for {
		var ws syscall.WaitStatus
		if _, err := syscall.Wait4(s.pid, &ws, syscall.WUNTRACED, nil); err != nil {
			return 0, err
		}
		if ws.Exited() {
			s.exited = true
			return 0, fmt.Errorf("debuggee exited with status %d", ws.ExitStatus())
		}
		if ws.Signaled() {
			s.exited = true
			return 0, fmt.Errorf("debuggee killed by signal %v", ws.Signal())
		}
		switch sig := ws.StopSignal(); sig {
		case syscall.SIGTRAP, syscall.SIGSTOP:
			s.stopped = true
			s.lastSignal = 0
			return sig, nil
		case syscall.SIGURG:
			// Go runtime asynchronous preemption: resume the tracee
			// without delivering the signal.
			s.lastSignal = 0
			if _, _, errno := syscall.Syscall6(
				syscall.SYS_PTRACE,
				uintptr(syscall.PTRACE_CONT),
				uintptr(s.pid),
				0, 0, 0, 0,
			); errno != 0 {
				return 0, fmt.Errorf("debug: PTRACE_CONT: %w", errno)
			}
		default:
			// A genuine signal-delivery-stop.  Report it; the caller
			// decides how to proceed.
			s.stopped = true
			s.lastSignal = sig
			return sig, nil
		}
	}
}

// LastSignal returns the signal of the most recent stop when that stop was
// a genuine signal-delivery-stop (a fault such as SIGSEGV, SIGFPE, SIGILL
// or SIGBUS), and 0 for breakpoint traps, single-steps, SIGSTOP and
// suppressed runtime signals.
func (s *Session) LastSignal() syscall.Signal { return s.lastSignal }

// Peek reads a word (8 bytes) from the debuggee's memory at addr.
func (s *Session) Peek(addr uint64) (uint64, error) {
	mem, err := os.OpenFile(fmt.Sprintf("/proc/%d/mem", s.pid), os.O_RDONLY, 0)
	if err != nil {
		return 0, fmt.Errorf("debug: open /proc/%d/mem: %w", s.pid, err)
	}
	defer mem.Close()
	buf := make([]byte, 8)
	if _, err := mem.ReadAt(buf, int64(addr)); err != nil {
		return 0, fmt.Errorf("debug: read mem %#x: %w", addr, err)
	}
	return uint64(buf[0]) | uint64(buf[1])<<8 | uint64(buf[2])<<16 | uint64(buf[3])<<24 |
		uint64(buf[4])<<32 | uint64(buf[5])<<40 | uint64(buf[6])<<48 | uint64(buf[7])<<56, nil
}

// Poke writes a word (8 bytes) to the debuggee's memory at addr.
func (s *Session) Poke(addr, val uint64) error {
	mem, err := os.OpenFile(fmt.Sprintf("/proc/%d/mem", s.pid), os.O_WRONLY, 0)
	if err != nil {
		return fmt.Errorf("debug: open /proc/%d/mem: %w", s.pid, err)
	}
	defer mem.Close()
	buf := []byte{byte(val), byte(val >> 8), byte(val >> 16), byte(val >> 24),
		byte(val >> 32), byte(val >> 40), byte(val >> 48), byte(val >> 56)}
	if _, err := mem.WriteAt(buf, int64(addr)); err != nil {
		return fmt.Errorf("debug: write mem %#x: %w", addr, err)
	}
	return nil
}

// ReadMemory reads len bytes from the debuggee's memory at addr.
func (s *Session) ReadMemory(addr uint64, length int) ([]byte, error) {
	out := make([]byte, length)
	for i := 0; i < length; i += 8 {
		word, err := s.Peek(addr + uint64(i))
		if err != nil {
			return out[:i], err
		}
		for j := 0; j < 8 && i+j < length; j++ {
			out[i+j] = byte(word >> (8 * j))
		}
	}
	return out, nil
}

// WriteMemory writes bytes to the debuggee's memory at addr.
func (s *Session) WriteMemory(addr uint64, data []byte) error {
	for i := 0; i < len(data); i += 8 {
		end := min(i+8, len(data))
		var word uint64
		for j := 0; j < end-i; j++ {
			word |= uint64(data[i+j]) << (8 * j)
		}
		if end-i < 8 {
			existing, err := s.Peek(addr + uint64(i))
			if err != nil {
				return err
			}
			mask := ^((uint64(1) << (8 * (end - i))) - 1)
			word = (existing & mask) | word
		}
		if err := s.Poke(addr+uint64(i), word); err != nil {
			return err
		}
	}
	return nil
}

// Step executes a single instruction in the debuggee.
func (s *Session) Step() error {
	if s.exited {
		return fmt.Errorf("debug: debuggee has exited")
	}
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_SINGLESTEP),
		uintptr(s.pid),
		0, 0, 0, 0,
	)
	if errno != 0 {
		return fmt.Errorf("debug: PTRACE_SINGLESTEP: %w", errno)
	}
	_, err := s.waitStopped()
	return err
}

// Continue resumes execution until the next breakpoint or exit.
func (s *Session) Continue() error {
	if s.exited {
		return fmt.Errorf("debug: debuggee has exited")
	}
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_CONT),
		uintptr(s.pid),
		0, 0, 0, 0,
	)
	if errno != 0 {
		return fmt.Errorf("debug: PTRACE_CONT: %w", errno)
	}
	_, err := s.waitStopped()
	return err
}

// Exited returns true if the debuggee has terminated.
func (s *Session) Exited() bool { return s.exited }

// Pid returns the debuggee's process ID.
func (s *Session) Pid() int { return s.pid }

// CodeBase returns the base address of the JIT code in the debuggee.
func (s *Session) CodeBase() uint64 { return s.codeBase }

// Kill terminates the debuggee and removes the session's scratch
// directory, so a successful session leaves no gasm-debug-* debris behind.
func (s *Session) Kill() {
	if !s.exited {
		syscall.Kill(s.pid, syscall.SIGKILL)
		syscall.Wait4(s.pid, nil, 0, nil)
		s.exited = true
	}
	if s.cmd != nil && s.cmd.Process != nil {
		s.cmd.Wait()
	}
	if s.tmpDir != "" {
		os.RemoveAll(s.tmpDir)
		s.tmpDir = ""
	}
}

// execRange is one executable mapping of the debuggee.
type execRange struct {
	lo, hi uint64
}

// execRanges parses the debuggee's executable mappings from /proc/pid/maps.
func execRanges(pid int) []execRange {
	data, err := os.ReadFile(fmt.Sprintf("/proc/%d/maps", pid))
	if err != nil {
		return nil
	}
	var out []execRange
	for line := range strings.SplitSeq(string(data), "\n") {
		fields := strings.Fields(line)
		if len(fields) < 2 || !strings.Contains(fields[1], "x") {
			continue
		}
		var lo, hi uint64
		if _, err := fmt.Sscanf(fields[0], "%x-%x", &lo, &hi); err == nil {
			out = append(out, execRange{lo, hi})
		}
	}
	return out
}

// findRWXMapping reads /proc/pid/maps and returns the base address of the
// first read-write-execute mapping (the JIT code region).
func findRWXMapping(pid int) uint64 {
	data, err := os.ReadFile(fmt.Sprintf("/proc/%d/maps", pid))
	if err != nil {
		return 0
	}
	for line := range strings.SplitSeq(string(data), "\n") {
		fields := strings.Fields(line)
		if len(fields) < 2 {
			continue
		}
		perms := fields[1]
		if len(perms) >= 3 && perms[0] == 'r' && perms[1] == 'w' && perms[2] == 'x' {
			var start uint64
			fmt.Sscanf(fields[0], "%x-", &start)
			return start
		}
	}
	return 0
}