#!/bin/bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "$0")/../.." && pwd)"
SESSION_NAME="${SESSION_NAME:-social-dev}"
COMPOSE_FILE="$ROOT_DIR/infra/docker/docker-compose.yml"
ENV_FILE="$ROOT_DIR/.env"
LITELLM_RUNTIME_CONFIG="$ROOT_DIR/.tmp/litellm-proxy-config.yaml"

usage() {
  echo "Usage: $0 {start|stop|restart}"
  echo ""
  echo "Commands:"
  echo "  start  Start LiteLLM + web + worker processes in tmux"
  echo "  stop   Stop tmux session and clean orphaned processes"
  echo "  restart Stop then start all app processes"
  exit 1
}

load_env_if_exists() {
  if [ -f "$ENV_FILE" ]; then
    set -a
    # shellcheck disable=SC1090
    . "$ENV_FILE"
    set +a
  fi
}

is_port_in_use() {
  local port="$1"

  if command -v lsof >/dev/null 2>&1; then
    lsof -iTCP:"$port" -sTCP:LISTEN -t >/dev/null 2>&1
    return $?
  fi

  if command -v ss >/dev/null 2>&1; then
    ss -ltn "sport = :$port" | awk 'NR > 1 {exit 0} END {exit 1}'
    return $?
  fi

  return 1
}

collect_listening_pids() {
  local port="$1"

  if command -v lsof >/dev/null 2>&1; then
    lsof -iTCP:"$port" -sTCP:LISTEN -t | sort -u
    return
  fi

  if command -v ss >/dev/null 2>&1; then
    ss -lptn "sport = :$port" | awk -F 'pid=' 'NF > 1 {split($2, tmp, ","); print tmp[1]}' | sort -u
  fi
}

kill_pids_gracefully() {
  local label="$1"
  shift
  local pids=("$@")
  local alive=()

  if [ "${#pids[@]}" -eq 0 ]; then
    return
  fi

  echo "Stopping ${label}: ${pids[*]}"
  kill -TERM "${pids[@]}" 2>/dev/null || true

  for _ in {1..10}; do
    alive=()
    for pid in "${pids[@]}"; do
      if kill -0 "$pid" 2>/dev/null; then
        alive+=("$pid")
      fi
    done

    if [ "${#alive[@]}" -eq 0 ]; then
      return
    fi

    sleep 1
  done

  echo "Force killing ${label}: ${alive[*]}"
  kill -KILL "${alive[@]}" 2>/dev/null || true
}

kill_matching_processes() {
  local label="$1"
  local pattern="$2"
  local pids

  pids="$(pgrep -f "$pattern" || true)"
  if [ -z "$pids" ]; then
    return
  fi

  # shellcheck disable=SC2086
  kill_pids_gracefully "$label" $pids
}

kill_listening_processes() {
  local label="$1"
  local port="$2"
  local pids

  pids="$(collect_listening_pids "$port" || true)"
  if [ -z "$pids" ]; then
    return
  fi

  # shellcheck disable=SC2086
  kill_pids_gracefully "$label" $pids
}

start() {
  echo "=== App Up ==="
  echo "This script starts LiteLLM + web + worker processes in tmux."
  echo "NOTE: Bootstrap (migrate + init-data) must be run separately."
  echo ""

  if ! command -v tmux >/dev/null 2>&1; then
    echo "Error: tmux is required." >&2
    exit 1
  fi

  if [ ! -f "$ENV_FILE" ]; then
    echo "Error: env file not found at $ENV_FILE" >&2
    exit 1
  fi

  if [ ! -f "$COMPOSE_FILE" ]; then
    echo "Error: compose file not found at $COMPOSE_FILE" >&2
    exit 1
  fi

  load_env_if_exists

  UVICORN_LOG_LEVEL="${SOCIAL_RUNTIME__LOG_LEVEL:-info}"
  UVICORN_LOG_LEVEL="$(echo "$UVICORN_LOG_LEVEL" | tr '[:upper:]' '[:lower:]')"
  WEB_PORT="${SOCIAL_WEB__PORT:-5775}"
  LITELLM_PORT="${SOCIAL_LITELLM__PORT:-3875}"

  if tmux has-session -t "$SESSION_NAME" 2>/dev/null; then
    echo "Error: tmux session '$SESSION_NAME' already exists." >&2
    echo "Hint: tmux kill-session -t $SESSION_NAME" >&2
    exit 1
  fi

  if is_port_in_use "$WEB_PORT"; then
    echo "Error: web port ${WEB_PORT} is already in use." >&2
    echo "Hint: run '$0 stop' or change SOCIAL_WEB__PORT in .env" >&2
    exit 1
  fi

  if is_port_in_use "$LITELLM_PORT"; then
    echo "Error: litellm port ${LITELLM_PORT} is already in use." >&2
    echo "Hint: run '$0 stop' or change SOCIAL_LITELLM__PORT in .env" >&2
    exit 1
  fi

  if [ -z "${SOCIAL_LLM__PROVIDER_KEYS__DASHSCOPE:-}" ]; then
    echo "Warning: SOCIAL_LLM__PROVIDER_KEYS__DASHSCOPE is empty; qwen calls may fail." >&2
  fi
  if [ -z "${SOCIAL_LLM__PROVIDER_KEYS__DEEPSEEK:-}" ]; then
    echo "Warning: SOCIAL_LLM__PROVIDER_KEYS__DEEPSEEK is empty; deepseek calls may fail." >&2
  fi

  echo "Starting LiteLLM + web + worker processes in tmux session '$SESSION_NAME'..."

  cd "$ROOT_DIR" && PYTHONPATH=backend/src uv run python backend/scripts/build_litellm_proxy_config.py --output "$LITELLM_RUNTIME_CONFIG"

  LITELLM_CMD="cd '$ROOT_DIR' && set -a && . '$ENV_FILE' && set +a && uv run litellm --config '$LITELLM_RUNTIME_CONFIG' --port ${LITELLM_PORT}"

  WEB_CMD="cd '$ROOT_DIR' && PYTHONPATH=backend/src SOCIAL_RUNTIME__SERVICE_NAME=web uv run uvicorn app:app --host \
${SOCIAL_WEB__HOST:-0.0.0.0} --port ${WEB_PORT} --workers \
${SOCIAL_WEB__WORKERS:-2} --log-level ${UVICORN_LOG_LEVEL}"

  WORKER_CRITICAL_CMD="cd '$ROOT_DIR' && PYTHONPATH=backend/src SOCIAL_RUNTIME__SERVICE_NAME=worker-critical uv run taskiq worker core.taskiq.app:critical_broker core.agentscope.runtime.tasks --workers ${SOCIAL_WORKER__GROUPS__CRITICAL__CONCURRENCY:-2}"
  WORKER_DEFAULT_CMD="cd '$ROOT_DIR' && PYTHONPATH=backend/src SOCIAL_RUNTIME__SERVICE_NAME=worker-default uv run taskiq worker core.taskiq.app:default_broker core.agentscope.runtime.tasks --workers ${SOCIAL_WORKER__GROUPS__DEFAULT__CONCURRENCY:-2}"
  WORKER_BULK_CMD="cd '$ROOT_DIR' && PYTHONPATH=backend/src SOCIAL_RUNTIME__SERVICE_NAME=worker-bulk uv run taskiq worker core.taskiq.app:bulk_broker core.agentscope.runtime.tasks --workers ${SOCIAL_WORKER__GROUPS__BULK__CONCURRENCY:-1}"

  tmux new-session -d -s "$SESSION_NAME" -n litellm "bash -lc \"$LITELLM_CMD; echo '[litellm] exited'; exec bash\""
  tmux new-window -t "$SESSION_NAME" -n web "bash -lc \"$WEB_CMD; echo '[web] exited'; exec bash\""
  tmux new-window -t "$SESSION_NAME" -n worker-critical "bash -lc \"$WORKER_CRITICAL_CMD; echo '[worker-critical] exited'; exec bash\""
  tmux new-window -t "$SESSION_NAME" -n worker-default "bash -lc \"$WORKER_DEFAULT_CMD; echo '[worker-default] exited'; exec bash\""
  tmux new-window -t "$SESSION_NAME" -n worker-bulk "bash -lc \"$WORKER_BULK_CMD; echo '[worker-bulk] exited'; exec bash\""

  echo ""
  echo "=== App Started ==="
  echo "Log files will be created in logs/ directory:"
  echo "  - web.log, web.error.log"
  echo "  - worker-critical.log, worker-critical.error.log"
  echo "  - worker-default.log, worker-default.error.log"
  echo "  - worker-bulk.log, worker-bulk.error.log"
  echo ""
  echo "tmux attach -t $SESSION_NAME"
  echo "tmux list-windows -t $SESSION_NAME"
}

stop() {
  echo "=== App Down ==="
  load_env_if_exists
  WEB_PORT="${SOCIAL_WEB__PORT:-5775}"
  LITELLM_PORT="${SOCIAL_LITELLM__PORT:-3875}"

  if tmux has-session -t "$SESSION_NAME" 2>/dev/null; then
    echo "Stopping tmux session '$SESSION_NAME'..."
    tmux kill-session -t "$SESSION_NAME"
  else
    echo "No tmux session '$SESSION_NAME' found."
  fi

  echo "Checking for orphaned processes..."

  kill_matching_processes "uvicorn" "uv run uvicorn app:app"
  kill_matching_processes "litellm" "uv run litellm --config"
  kill_matching_processes "taskiq workers" "uv run taskiq worker core.taskiq.app:"

  kill_listening_processes "port ${WEB_PORT} listeners" "$WEB_PORT"
  kill_listening_processes "port ${LITELLM_PORT} listeners" "$LITELLM_PORT"

  if is_port_in_use "$WEB_PORT"; then
    echo "Warning: port ${WEB_PORT} is still in use after cleanup." >&2
    echo "Hint: check process with 'lsof -iTCP:${WEB_PORT} -sTCP:LISTEN'" >&2
    return 1
  fi

  if is_port_in_use "$LITELLM_PORT"; then
    echo "Warning: port ${LITELLM_PORT} is still in use after cleanup." >&2
    echo "Hint: check process with 'lsof -iTCP:${LITELLM_PORT} -sTCP:LISTEN'" >&2
    return 1
  fi

  echo "Session stopped and cleaned up."
}

restart() {
  stop
  echo ""
  start
}

case "${1:-}" in
  start) start ;;
  stop) stop ;;
  restart) restart ;;
  *) usage ;;
esac