#!/bin/bash
# debug_guest_init.sh - One-command debug launch for guest NVIDIA driver initialization
# Usage: ./debug_guest_init.sh [options]

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
BASE_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"

# Defaults
GDB_PORT=1234
FREEZE_AT_START=true
MONITOR_GUEST=true
GUEST_SERIAL="/dev/ttyS0"
DURATION=300
OUTPUT_DIR="/tmp/gsp_emu_debug_$(date +%Y%m%d_%H%M%S)"

# Colors
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_warn() { echo -e "${YELLOW}[WARN]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1"; }
log_ok() { echo -e "${GREEN}[OK]${NC} $1"; }

usage() {
    cat <<EOF
Usage: $0 [options]

Debug guest NVIDIA driver initialization with full observability.

Options:
  --gdb-port PORT         GDB server port (default: 1234)
  --no-freeze             Don't freeze QEMU at startup (-S flag)
  --no-monitor            Disable guest init monitor
  --guest-serial DEV      Guest serial device (default: /dev/ttyS0)
  --duration SEC          Monitor duration in seconds (default: 300)
  --output-dir DIR        Output directory (default: /tmp/gsp_emu_debug_TIMESTAMP)
  --vmlinux PATH          Path to guest vmlinux with debug symbols
  --help                  Show this help

This script:
  1. Enables kernel tracepoints for GSP emulation
  2. Launches QEMU with GDB server
  3. Starts guest init monitor (correlates host traces + guest serial)
  4. Opens GDB with nvidia_init.gdb script
  5. Saves all logs to output directory

Example:
  $0 --vmlinux /path/to/guest/vmlinux
  $0 --no-freeze --duration 600
EOF
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --gdb-port) GDB_PORT="$2"; shift 2 ;;
        --no-freeze) FREEZE_AT_START=false; shift ;;
        --no-monitor) MONITOR_GUEST=false; shift ;;
        --guest-serial) GUEST_SERIAL="$2"; shift 2 ;;
        --duration) DURATION="$2"; shift 2 ;;
        --output-dir) OUTPUT_DIR="$2"; shift 2 ;;
        --vmlinux) VMLINUX="$2"; shift 2 ;;
        --help) usage; exit 0 ;;
        *) log_error "Unknown option: $1"; usage; exit 1 ;;
    esac
done

mkdir -p "$OUTPUT_DIR"
log_info "Output directory: $OUTPUT_DIR"

# Check prerequisites
log_info "Checking prerequisites..."

# Check for mdev
MDEV_UUID=$(ls /sys/bus/mdev/devices/ 2>/dev/null | head -1)
if [[ -z "$MDEV_UUID" ]]; then
    log_error "No mdev device found. Run create_mdev.sh first."
    exit 1
fi
log_info "Using mdev: $MDEV_UUID"

# Build QEMU command
QEMU_CMD=(
    qemu-system-x86_64
    -enable-kvm
    -cpu host
    -smp 4
    -m 8G
    -device vfio-pci,sysfsdev=/sys/bus/mdev/devices/$MDEV_UUID,display=off
    -drive file=/var/lib/libvirt/images/ubuntu-vgpu.qcow2,format=qcow2,if=virtio
    -netdev user,id=net0,hostfwd=tcp::2222-:22
    -device virtio-net-pci,netdev=net0
    -nographic
    -serial mon:stdio
)

if [[ "$FREEZE_AT_START" == "true" ]]; then
    QEMU_CMD+=(-s -S)
else
    QEMU_CMD+=(-s)
fi

if [[ "$GDB_PORT" != "1234" ]]; then
    QEMU_CMD+=(-gdb "tcp::$GDB_PORT")
fi

# Start guest init monitor in background
MONITOR_PID=""
if [[ "$MONITOR_GUEST" == "true" ]]; then
    log_info "Starting guest init monitor..."
    python3 "$BASE_DIR/tools/guest_init_monitor.py" \
        --guest-serial "$GUEST_SERIAL" \
        --duration "$DURATION" \
        > "$OUTPUT_DIR/guest_monitor.log" 2>&1 &
    MONITOR_PID=$!
    log_ok "Guest monitor started (PID: $MONITOR_PID)"
fi

# Start host trace capture
log_info "Starting host trace capture..."
sudo cat /sys/kernel/debug/tracing/trace_pipe > "$OUTPUT_DIR/host_trace.log" 2>&1 &
TRACE_PID=$!
log_ok "Host trace capture started (PID: $TRACE_PID)"

# Save QEMU command
echo "${QEMU_CMD[@]}" > "$OUTPUT_DIR/qemu_cmd.txt"

# Launch QEMU in background
log_info "Launching QEMU with GDB server on port $GDB_PORT..."
"${QEMU_CMD[@]}" > "$OUTPUT_DIR/qemu.log" 2>&1 &
QEMU_PID=$!
log_ok "QEMU started (PID: $QEMU_PID)"

# Give QEMU time to start
sleep 3

# Launch GDB if vmlinux provided
if [[ -n "${VMLINUX:-}" && -f "$VMLINUX" ]]; then
    log_info "Launching GDB with nvidia_init.gdb..."
    log_info "GDB will connect to localhost:$GDB_PORT"
    log_info "In GDB: (gdb) target remote localhost:$GDB_PORT"
    log_info "Then: (gdb) continue"
    
    gdb -x "$BASE_DIR/tools/gdb_scripts/nvidia_init.gdb" "$VMLINUX" \
        -ex "target remote localhost:$GDB_PORT" \
        -ex "set logging file $OUTPUT_DIR/gdb.log" \
        -ex "set logging on"
else
    log_warn "No vmlinux provided. Connect manually:"
    log_warn "  gdb -x $BASE_DIR/tools/gdb_scripts/nvidia_init.gdb /path/to/vmlinux"
    log_warn "  (gdb) target remote localhost:$GDB_PORT"
    log_warn ""
    log_warn "Press Ctrl+C to stop QEMU and monitoring..."
    
    # Wait for QEMU or interrupt
    wait $QEMU_PID
fi

# Cleanup
log_info "Cleaning up..."
[[ -n "$MONITOR_PID" ]] && kill $MONITOR_PID 2>/dev/null || true
[[ -n "$TRACE_PID" ]] && kill $TRACE_PID 2>/dev/null || true
kill $QEMU_PID 2>/dev/null || true

log_info "Debug session complete. Logs in: $OUTPUT_DIR"
ls -la "$OUTPUT_DIR"