#!/bin/bash

set -euo pipefail

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

if ! command -v aws >/dev/null 2>&1; then
  echo "AWS CLI not found. Installing via install_awscli.sh..."
  "${SCRIPT_DIR}/install_awscli.sh"
  echo "AWS CLI installed."
else
  echo "AWS CLI already available."
fi

if ! command -v uvx >/dev/null 2>&1; then
  echo "uvx command is required but not found. Please install uv and retry." >&2
  exit 1
fi

APT_UPDATED=0
ensure_command() {
  local cmd="$1"
  local pkg="$2"

  if command -v "${cmd}" >/dev/null 2>&1; then
    return
  fi

  if ! command -v apt-get >/dev/null 2>&1; then
    echo "${cmd} is required but apt-get is unavailable to install ${pkg}." >&2
    exit 1
  fi

  if [ "${APT_UPDATED}" -eq 0 ]; then
    echo "Updating package index to install missing dependencies..."
    apt-get update -qq >/dev/null
    APT_UPDATED=1
  fi

  echo "Installing ${pkg} to provide ${cmd}..."
  apt-get install -y -qq "${pkg}" >/dev/null
  echo "${cmd} installed."
}

ensure_command ps procps
ensure_command grep grep
ensure_command awk gawk

mapfile -t PIDS < <(ps aux | grep gunicorn | grep -v grep | awk '{print $2}')

if [ ${#PIDS[@]} -eq 0 ]; then
  echo "No gunicorn processes found. Nothing to dump."
  exit 0
fi

OUTPUT_DIR="${SCRIPT_DIR}/thread_dumps"
mkdir -p "${OUTPUT_DIR}"

TIMESTAMP="$(date -u +%Y-%m-%dT%H%M%SZ)"
HOSTNAME="$(hostname)"

for PID in "${PIDS[@]}"; do
  FILE_NAME="${TIMESTAMP}_${HOSTNAME}_pid${PID}.txt"
  FILE_PATH="${OUTPUT_DIR}/${FILE_NAME}"
  echo "Collecting thread dump for PID ${PID}..."
  if uvx py-spy dump --pid "${PID}" > "${FILE_PATH}"; then
    echo "Thread dump saved to ${FILE_PATH}"
    echo "Uploading ${FILE_NAME} to S3..."
    aws s3 cp "${FILE_PATH}" "s3://suno-debug/thread-dumps/studio-api/${FILE_NAME}" --quiet
    echo "Uploaded to s3://suno-debug/thread-dumps/studio-api/${FILE_NAME}"
  else
    echo "Failed to collect thread dump for PID ${PID}" >&2
  fi
done

echo "Thread dump collection complete."