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

//go:build linux

package main

import (
	"fmt"
	"io"
	"os"
	"sort"
	"strings"
	"time"

	"sourcedock.dev/petrbalvin/gasm-devkit/debug"
	"sourcedock.dev/petrbalvin/gasm-devkit/verify"
)

func cmdDebug(args []string) int {
	fs := newCommand("debug", "gasm debug <file.s> --func <name>", `
Interactive debugger for JIT-assembled functions.  Launches the
function in a traced subprocess (ptrace), then provides a REPL for
single-stepping, breakpoints, register and memory inspection.

REPL commands:
  break <label|addr> [if <reg> <op> <val>]
                       set a breakpoint, optionally conditional on a
                       register comparison (reg-reg or reg-immediate)
  delete <label|addr>  remove a breakpoint
  info break           list all breakpoints
  step [n], s          single-step n instructions (default 1)
  next, n              step over a CALL
  finish, fin          run until the function returns
  continue, c          run until a breakpoint, watchpoint or exit
  disas [n], u         disassemble n instructions at PC
  regs                 print general-purpose and vector registers
  where                show source line and nearest label at PC
  stack                show stack near RSP (return address + ABI0 args)
  bt, backtrace        backtrace (current frame + return address)
  x [addr] [len]       hex-dump memory (default: current PC, 64 bytes)
  w <addr> <val...>    write bytes to memory
  set <reg> <value>    set a register
  watch <addr> [r|w] [size]
                       set a hardware watchpoint (write by default)
  unwatch [<slot>]     clear one watchpoint, or all without an argument
  labels, l            list function labels and offsets
  help, h, ?           show command help
  quit, q              kill the debuggee and exit
`)
	funcName := fs.String("func", "", "function to debug")
	argsFile := fs.String("args", "", "file containing the ABI0 argument block")
	bufSpec := fs.String("buf", "", "buffer specification: name:size:pattern[,name:size:pattern...] where pattern is zero, ones, seq, or hex")
	script := fs.String("script", "", "run REPL commands from a file (one per line) and exit; '-' reads stdin")
	cover := fs.Bool("cover", false, "run to completion with a breakpoint on every instruction and report which executed and how often")
	timeout := fs.Duration("timeout", 0, "kill the debuggee after this duration (e.g. 30s); for headless --script runs")
	fs.Parse(args)

	// --- Debuggee mode (internal, spawned by the debugger) ---
	if os.Getenv("GASM_DEBUG_TARGET") != "" {
		tmpDir := os.Getenv("GASM_DEBUG_TMP")
		if tmpDir == "" || fs.NArg() < 1 || *funcName == "" || *argsFile == "" {
			fmt.Fprintln(os.Stderr, "gasm debug: internal debuggee mode")
			return 2
		}
		if err := debug.RunTarget(fs.Arg(0), *funcName, *argsFile, tmpDir); err != nil {
			fmt.Fprintf(os.Stderr, "gasm debug: %v\n", err)
			return 1
		}
		return 0
	}

	// --- Debugger mode (interactive REPL) ---
	if fs.NArg() < 1 || *funcName == "" {
		fmt.Fprintln(os.Stderr, "usage: gasm debug <file.s> --func <name>")
		return 2
	}
	path := fs.Arg(0)

	// The watchdog is armed before anything can block: ptrace attach and a
	// continued kernel loop both hang the run when the environment forbids
	// tracing or the kernel loops forever, and neither is interruptible from
	// the inside.
	if *timeout > 0 {
		go func() {
			time.Sleep(*timeout)
			fmt.Fprintf(os.Stderr, "gasm debug: timeout (%s), killing the debuggee\n", *timeout)
			os.Exit(3)
		}()
	}

	// Load the kernel to extract function metadata and labels.
	k, err := verify.Load(path)
	if err != nil {
		fmt.Fprintf(os.Stderr, "gasm debug: %v\n", err)
		return 1
	}
	defer k.Close()

	fl, err := k.Func(*funcName)
	if err != nil {
		fmt.Fprintf(os.Stderr, "gasm debug: %v\n", err)
		return 1
	}

	// Build the label list for the REPL.
	var labels []debug.Label
	for name, off := range fl.Labels {
		labels = append(labels, debug.Label{Name: name, Offset: off})
	}
	sort.Slice(labels, func(i, j int) bool { return labels[i].Offset < labels[j].Offset })

	// Launch the debuggee with the argument block.
	var argBlock []byte
	var bufAddrs []uint64
	var sess *debug.Session
	if *bufSpec != "" {
		// Parse the function signature to determine argument layout.
		src, err := readSource(path)
		if err != nil {
			fmt.Fprintf(os.Stderr, "gasm debug: %v\n", err)
			return 1
		}
		sig, ok := verify.ExtractFuncSig(src, *funcName)
		if !ok {
			fmt.Fprintf(os.Stderr, "gasm debug: no // func signature found for %s\n", *funcName)
			return 1
		}
		layout := verify.ArgLayout(sig)

		// Parse the buffer spec to get buffer names.
		bufNames := parseBufNames(*bufSpec)

		// Allocate buffers in the debuggee.
		argBlock = make([]byte, fl.Args)
		sess, bufAddrs, err = debug.LaunchWithBuffers("", path, *funcName, argBlock, *bufSpec)
		if err != nil {
			fmt.Fprintf(os.Stderr, "gasm debug: %v\n", err)
			return 1
		}

		// Construct the argument block with buffer pointers at the correct positions.
		for _, arg := range layout {
			if !arg.IsPtr {
				continue
			}
			// Find the buffer that matches this argument.
			for i, name := range bufNames {
				if i < len(bufAddrs) && (name == arg.Name || strings.HasPrefix(arg.Name, name)) {
					addr := bufAddrs[i]
					off := arg.Offset
					if off+8 <= len(argBlock) {
						argBlock[off] = byte(addr)
						argBlock[off+1] = byte(addr >> 8)
						argBlock[off+2] = byte(addr >> 16)
						argBlock[off+3] = byte(addr >> 24)
						argBlock[off+4] = byte(addr >> 32)
						argBlock[off+5] = byte(addr >> 40)
						argBlock[off+6] = byte(addr >> 48)
						argBlock[off+7] = byte(addr >> 56)
					}
					// For slices, also set the length and capacity.
					if strings.HasPrefix(arg.Typ, "[]") && off+24 <= len(argBlock) {
						// Find the buffer size from the spec.
						size := parseBufSize(*bufSpec, name)
						// Length at offset+8, capacity at offset+16.
						for j := range 8 {
							argBlock[off+8+j] = byte(size >> (j * 8))
							argBlock[off+16+j] = byte(size >> (j * 8))
						}
					}
					break
				}
			}
		}
	} else {
		argBlock = make([]byte, fl.Args)
		sess, err = debug.Launch("", path, *funcName, argBlock)
		if err != nil {
			fmt.Fprintf(os.Stderr, "gasm debug: %v\n", err)
			return 1
		}
	}
	defer sess.Kill()

	bm := debug.NewBreakpoints(sess)
	fmt.Printf("gasm debug: %s in %s (pid %d)\n", *funcName, path, sess.Pid())

	// Convert the line table for the command loop.
	var srcLines []debug.SourceLine
	for _, le := range fl.Lines {
		srcLines = append(srcLines, debug.SourceLine{Offset: le.Offset, Line: le.Line})
	}

	// Coverage mode: pre-register a breakpoint on every instruction (walked
	// by length through the function body while the debuggee is stopped) and
	// let the kernel run to completion.  Each trap counts a hit for that
	// instruction, so the final report shows exactly which instructions
	// executed and how often, with the label-level view derived from it.
	// Expect the run to slow to ptrace speed: one trap per executed
	// instruction.
	if *cover {
		base := sess.CodeBase() + uint64(fl.Offset)
		type coverInstr struct {
			off  uint64
			text string
		}
		var instrs []coverInstr
		for off := uint64(0); off < uint64(fl.Size); {
			text, ln, err := sess.Disassemble(base + off)
			if err != nil || ln == 0 {
				break
			}
			instrs = append(instrs, coverInstr{off: off, text: text})
			off += uint64(ln)
		}
		for _, in := range instrs {
			if _, err := bm.SetWithCond(base+in.off, fmt.Sprintf("func+%#x", in.off), nil); err != nil {
				fmt.Fprintf(os.Stderr, "gasm debug: cover: %v\n", err)
				return 1
			}
		}
		fmt.Printf("gasm debug: coverage run over %d instructions\n", len(instrs))
		for {
			for _, bp := range bm.All() {
				bm.Reinsert(bp.Addr)
			}
			if err := sess.Continue(); err != nil {
				break // debuggee finished or died
			}
			if sess.Exited() {
				break
			}
			regs, rerr := sess.GetRegs()
			if rerr != nil {
				break
			}
			// HandleTrap restores the original byte, rewinds PC and counts
			// the hit on the breakpoint itself.  Single-step over the
			// restored instruction so the reinsertion at the top of the
			// loop cannot re-trap on the same breakpoint.
			if bp := bm.HandleTrap(&regs); bp != nil {
				if err := sess.Step(); err != nil {
					break
				}
			}
		}
		hits := map[uint64]int{}
		traps := 0
		for _, bp := range bm.All() {
			if n := bp.Hits(); n > 0 {
				hits[bp.Addr-base] = n
				traps += n
			}
		}
		var hit []string
		var missed []string
		for _, l := range labels {
			if hits[uint64(l.Offset)] > 0 {
				hit = append(hit, l.Name)
			} else {
				missed = append(missed, l.Name)
			}
		}
		sort.Strings(hit)
		sort.Strings(missed)
		fmt.Printf("coverage: %d/%d instructions executed (%d traps)\n", len(hits), len(instrs), traps)
		fmt.Printf("coverage: %d/%d labels reached\n", len(hit), len(labels))
		for _, l := range hit {
			fmt.Printf("  covered   %s\n", l)
		}
		for _, l := range missed {
			fmt.Printf("  MISSED    %s\n", l)
		}
		fmt.Println("executed instructions:")
		for _, in := range instrs {
			if n := hits[in.off]; n > 0 {
				fmt.Printf("  func+%#04x  %4dx  %s\n", in.off, n, in.text)
			}
		}
		return 0
	}

	// Headless mode: run the script through the normal command loop and
	// exit.  The watchdog armed above covers launch, continue and step.
	var in io.Reader = os.Stdin
	if *script != "" {
		if *script == "-" {
			in = os.Stdin
		} else {
			f, err := os.Open(*script)
			if err != nil {
				fmt.Fprintf(os.Stderr, "gasm debug: %v\n", err)
				return 1
			}
			defer f.Close()
			in = f
		}
	}
	debug.REPL(sess, bm, sess.CodeBase(), fl.Offset, fl.Size, fl.Args, labels, srcLines, in)
	return 0
}

// parseBufNames extracts buffer names from a buffer specification.
// Format: name:size:pattern[,name:size:pattern...]
func parseBufNames(spec string) []string {
	var names []string
	for part := range strings.SplitSeq(spec, ",") {
		fields := strings.SplitN(part, ":", 3)
		if len(fields) >= 1 && fields[0] != "" {
			names = append(names, fields[0])
		}
	}
	return names
}

// parseBufSize extracts the size of a named buffer from a buffer specification.
func parseBufSize(spec, name string) int {
	for part := range strings.SplitSeq(spec, ",") {
		fields := strings.SplitN(part, ":", 3)
		if len(fields) >= 2 && fields[0] == name {
			var size int
			fmt.Sscanf(fields[1], "%d", &size)
			return size
		}
	}
	return 0
}