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

// 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.
func LaunchWithBuffers(gasmBin, asmPath, funcName string, args []byte, bufSpec string) (*Session, []uint64, error) {
	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", "--target", "--func", funcName, "--args", argsFile, asmPath)
	cmd.Env = append(os.Environ(), "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}

	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, nil, fmt.Errorf("debug: wait for stop: %w", err)
	}

	entryFile := filepath.Join(tmpDir, "entry")
	for i := 0; i < 500; i++ {
		if _, err := os.Stat(entryFile); err == nil {
			break
		}
		time.Sleep(5 * time.Millisecond)
	}

	if err := s.Continue(); err != nil {
		return nil, nil, fmt.Errorf("debug: continue to entry: %w", err)
	}

	if _, err := syscall.Wait4(s.pid, &ws, syscall.WUNTRACED, nil); err != nil {
		return nil, nil, fmt.Errorf("debug: wait for entry: %w", err)
	}

	s.stopped = true

	s.codeBase = findRWXMapping(s.pid)
	if s.codeBase == 0 {
		baseFile := filepath.Join(tmpDir, "codebase")
		if data, err := os.ReadFile(baseFile); err == nil {
			fmt.Sscanf(string(data), "%d", &s.codeBase)
		}
	}

	var bufAddrs []uint64
	if bufSpec != "" {
		addrFile := filepath.Join(tmpDir, "bufaddrs")
		if data, err := os.ReadFile(addrFile); err == nil {
			for _, line := range strings.Split(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
}

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

// 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 := 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)
		}
		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)
	}
	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") {
		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
}