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

//go:build linux && amd64

// Package debug implements the interactive debugger for gasm (Phase 4):
// single-stepping, breakpoints, register and memory inspection for
// JIT-assembled Plan 9 amd64 functions, controlled via ptrace.
package debug

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

// Regs holds the full general-purpose register set of a traced process
// (the Linux amd64 user_regs_struct layout).
type Regs struct {
	R15     uint64
	R14     uint64
	R13     uint64
	R12     uint64
	RBP     uint64
	RBX     uint64
	R11     uint64
	R10     uint64
	R9      uint64
	R8      uint64
	RAX     uint64
	RCX     uint64
	RDX     uint64
	RSI     uint64
	RDI     uint64
	OrigRAX uint64
	RIP     uint64
	CS      uint64
	RFLAGS  uint64
	RSP     uint64
	SS      uint64
	FSBase  uint64
	GSBase  uint64
	DS      uint64
	ES      uint64
	FS      uint64
	GS      uint64
}

// 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
}

// Launch starts the debuggee subprocess (gasm debug --target ...) and
// attaches to it via ptrace.  The debuggee assembles the file, maps the
// JIT code, calls PTRACE_TRACEME and raises SIGSTOP; Launch waits for
// that initial stop and returns a ready Session.
func Launch(gasmBin, asmPath, funcName string, args []byte) (*Session, error) {
	self, err := os.Executable()
	if err != nil {
		return nil, fmt.Errorf("debug: cannot find gasm binary: %w", err)
	}
	if gasmBin != "" {
		self = gasmBin
	}

	// Write the arg block to a temp file (the child reads it).
	tmpDir, err := os.MkdirTemp("", "gasm-debug-*")
	if err != nil {
		return 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, fmt.Errorf("debug: write args: %w", err)
	}

	cmd := exec.Command(self, "debug", "--target", "--func", funcName, "--args", argsFile, asmPath)
	cmd.Env = append(os.Environ(), "GASM_DEBUG_TMP="+tmpDir)
	cmd.Stdout = nil // output goes to the debugger, not the terminal
	cmd.Stderr = os.Stderr
	cmd.SysProcAttr = &syscall.SysProcAttr{}

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

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

	// Wait for the child to signal readiness and stop.  The child calls
	// PTRACE_TRACEME then SIGSTOP, so Wait4 with WUNTRACED observes the
	// ptrace-stop directly (no PTRACE_ATTACH needed).
	readyFile := filepath.Join(tmpDir, "ready")
	for i := 0; i < 500; i++ {
		if _, err := os.Stat(readyFile); err == nil {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}
	var ws syscall.WaitStatus
	if _, err := syscall.Wait4(s.pid, &ws, syscall.WUNTRACED, nil); err != nil {
		cmd.Process.Kill()
		os.RemoveAll(tmpDir)
		return nil, fmt.Errorf("debug: wait for debuggee: %w", err)
	}
	s.stopped = true

	// Read the code base from /proc/pid/maps (find the RWX mapping).
	s.codeBase = findRWXMapping(s.pid)
	if s.codeBase == 0 {
		// Fallback: try the file the child wrote.
		baseFile := filepath.Join(tmpDir, "codebase")
		if data, err := os.ReadFile(baseFile); err == nil {
			fmt.Sscanf(string(data), "%d", &s.codeBase)
		}
	}

	return s, nil
}

// wait waits for the debuggee to stop and returns the wait status.
func (s *Session) wait() error {
	var ws syscall.WaitStatus
	_, err := syscall.Wait4(s.pid, &ws, 0, nil)
	if err != nil {
		return err
	}
	if ws.Exited() {
		s.exited = true
		return fmt.Errorf("debuggee exited with status %d", ws.ExitStatus())
	}
	s.stopped = true
	return nil
}

// GetRegs reads the general-purpose registers of the stopped debuggee.
func (s *Session) GetRegs() (Regs, error) {
	var regs Regs
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_GETREGS),
		uintptr(s.pid),
		0,
		uintptr(unsafe.Pointer(&regs)),
		0, 0,
	)
	if errno != 0 {
		return regs, fmt.Errorf("debug: PTRACE_GETREGS: %w", errno)
	}
	return regs, nil
}

// SetRegs writes the general-purpose registers of the stopped debuggee.
func (s *Session) SetRegs(regs *Regs) error {
	_, _, errno := syscall.Syscall6(
		syscall.SYS_PTRACE,
		uintptr(syscall.PTRACE_SETREGS),
		uintptr(s.pid),
		0,
		uintptr(unsafe.Pointer(regs)),
		0, 0,
	)
	if errno != 0 {
		return fmt.Errorf("debug: PTRACE_SETREGS: %w", errno)
	}
	return nil
}

// Peek reads a word (8 bytes) from the debuggee's memory at addr.
// Uses /proc/pid/mem which works reliably with Go's multi-threaded runtime.
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 := i + 8
		if end > len(data) {
			end = len(data)
		}
		var word uint64
		for j := 0; j < end-i; j++ {
			word |= uint64(data[i+j]) << (8 * j)
		}
		// For partial writes, read-modify-write the existing word.
		if end-i < 8 {
			existing, err := s.Peek(addr + uint64(i))
			if err != nil {
				return err
			}
			// Clear the bytes we're overwriting and merge.
			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)
	}
	return s.wait()
}

// 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)
	}
	return s.wait()
}

// 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.
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()
	}
}

// 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.Split(string(data), "\n") {
		// Format: addr-addr perms offset dev inode pathname
		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' {
			// Parse the start address.
			var start uint64
			fmt.Sscanf(fields[0], "%x-", &start)
			return start
		}
	}
	return 0
}