#!/usr/bin/env bash
# =============================================================================
# DocEngine — Backend Startup Script
# =============================================================================
# Usage:
#   ./run_docengine.sh              Run the API server (default)
#   ./run_docengine.sh server       Run the API server only
#   ./run_docengine.sh worker       Run the Celery worker only
#   ./run_docengine.sh all          Run both API server and Celery worker
#   ./run_docengine.sh setup        Install deps + run migrations only
#   ./run_docengine.sh migrate      Run Alembic migrations only
# =============================================================================

set -euo pipefail

# ── Paths ────────────────────────────────────────────────────────────────────
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_DIR="${SCRIPT_DIR}/docengine"
VENV_DIR="${PROJECT_DIR}/.venv"
PID_DIR="${PROJECT_DIR}/.pids"

# ── Colors ───────────────────────────────────────────────────────────────────
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

log()   { echo -e "${GREEN}[DocEngine]${NC} $*"; }
warn()  { echo -e "${YELLOW}[DocEngine]${NC} $*"; }
err()   { echo -e "${RED}[DocEngine]${NC} $*" >&2; }

# ── Pre-flight checks ───────────────────────────────────────────────────────
check_python() {
    if command -v python3 &>/dev/null; then
        PYTHON=python3
    elif command -v python &>/dev/null; then
        PYTHON=python
    else
        err "Python 3 is not installed. Please install Python 3.12+."
        exit 1
    fi

    local version
    version=$($PYTHON -c 'import sys; print(f"{sys.version_info.major}.{sys.version_info.minor}")')
    log "Using Python ${version} (${PYTHON})"
}

check_redis() {
    if command -v redis-cli &>/dev/null; then
        if redis-cli ping &>/dev/null; then
            log "Redis is running ✓"
        else
            warn "Redis is installed but not running. Celery worker will fail without Redis."
            warn "Start Redis with: brew services start redis  (or)  redis-server --daemonize yes"
        fi
    else
        warn "redis-cli not found. Celery worker requires Redis."
    fi
}

# ── Virtual Environment ─────────────────────────────────────────────────────
ensure_venv() {
    local created=false
    if [ ! -d "${VENV_DIR}" ]; then
        log "Creating virtual environment at ${VENV_DIR}..."
        $PYTHON -m venv "${VENV_DIR}"
        log "Virtual environment created ✓"
        created=true
    fi
    # shellcheck disable=SC1091
    source "${VENV_DIR}/bin/activate"
    log "Virtual environment activated ✓"

    # Auto-install dependencies if this is a fresh venv or if uvicorn is missing
    if [ "${created}" = true ] || ! python -c "import uvicorn" &>/dev/null; then
        log "Dependencies missing. Auto-installing dependencies..."
        install_deps
    fi
}

# ── Dependencies ─────────────────────────────────────────────────────────────
install_deps() {
    log "Installing Python dependencies..."
    pip install --upgrade pip --quiet
    pip install -r "${PROJECT_DIR}/requirements.txt" --quiet
    log "Dependencies installed ✓"
}

# ── Storage Directories ─────────────────────────────────────────────────────
ensure_storage() {
    local dirs=("documents" "templates" "images" "temp" "rendered")
    for dir in "${dirs[@]}"; do
        mkdir -p "${PROJECT_DIR}/storage/${dir}"
    done
    log "Storage directories ready ✓"
}

# ── Database Migrations ─────────────────────────────────────────────────────
run_migrations() {
    log "Running Alembic database migrations..."
    cd "${PROJECT_DIR}"
    if alembic upgrade head 2>/dev/null; then
        log "Migrations applied ✓"
    else
        warn "Alembic migrations failed or no migrations to apply."
        warn "If this is a fresh database, ensure PostgreSQL is reachable at the configured host."
    fi
}

# ── PID Management ───────────────────────────────────────────────────────────
mkdir_pids() {
    mkdir -p "${PID_DIR}"
}

save_pid() {
    echo "$2" > "${PID_DIR}/$1.pid"
}

remove_pid() {
    rm -f "${PID_DIR}/$1.pid"
}

read_pid() {
    local pidfile="${PID_DIR}/$1.pid"
    if [ -f "${pidfile}" ]; then
        cat "${pidfile}"
    fi
}

is_running() {
    local pid
    pid=$(read_pid "$1")
    if [ -n "${pid}" ] && kill -0 "${pid}" 2>/dev/null; then
        return 0
    fi
    return 1
}

# ── Cleanup on exit ─────────────────────────────────────────────────────────
cleanup() {
    log "Shutting down..."
    local server_pid worker_pid
    server_pid=$(read_pid "server")
    worker_pid=$(read_pid "worker")

    if [ -n "${server_pid}" ] && kill -0 "${server_pid}" 2>/dev/null; then
        log "Stopping API server (PID ${server_pid})..."
        kill "${server_pid}" 2>/dev/null || true
        remove_pid "server"
    fi

    if [ -n "${worker_pid}" ] && kill -0 "${worker_pid}" 2>/dev/null; then
        log "Stopping Celery worker (PID ${worker_pid})..."
        kill "${worker_pid}" 2>/dev/null || true
        remove_pid "worker"
    fi

    log "Shutdown complete."
}

# ── Start API Server ────────────────────────────────────────────────────────
start_server() {
    if is_running "server"; then
        warn "API server is already running (PID $(read_pid 'server'))"
        return
    fi

    cd "${PROJECT_DIR}"
    log "Starting API server on port 7989..."
    log "  Docs:   ${CYAN}http://localhost:7989/docs${NC}"
    log "  ReDoc:  ${CYAN}http://localhost:7989/redoc${NC}"
    log "  Health: ${CYAN}http://localhost:7989/api/v1/health${NC}"
    echo ""

    python -m uvicorn app.main:app \
        --host 0.0.0.0 \
        --port 7989 \
        --reload \
        --log-level info &

    local pid=$!
    save_pid "server" "${pid}"
    log "API server started (PID ${pid}) ✓"
}

# ── Start Celery Worker ─────────────────────────────────────────────────────
start_worker() {
    if is_running "worker"; then
        warn "Celery worker is already running (PID $(read_pid 'worker'))"
        return
    fi

    cd "${PROJECT_DIR}"
    log "Starting Celery worker..."

    celery -A app.workers.celery_app worker \
        --loglevel=info \
        --concurrency=4 \
        --pool=prefork \
        -Q celery,document_processing \
        -E &

    local pid=$!
    save_pid "worker" "${pid}"
    log "Celery worker started (PID ${pid}) ✓"
}

# ── Setup (no server start) ─────────────────────────────────────────────────
do_setup() {
    check_python
    ensure_venv
    install_deps
    ensure_storage
    check_redis
    run_migrations
    log ""
    log "Setup complete. Run ${CYAN}./run_docengine.sh${NC} to start the server."
}

# ── Main ─────────────────────────────────────────────────────────────────────
main() {
    local command="${1:-server}"

    case "${command}" in
        setup)
            do_setup
            ;;
        migrate)
            check_python
            ensure_venv
            run_migrations
            ;;
        server)
            check_python
            ensure_venv
            ensure_storage
            mkdir_pids
            trap cleanup EXIT INT TERM
            start_server
            wait
            ;;
        worker)
            check_python
            ensure_venv
            check_redis
            mkdir_pids
            trap cleanup EXIT INT TERM
            start_worker
            wait
            ;;
        all)
            check_python
            ensure_venv
            ensure_storage
            check_redis
            mkdir_pids
            trap cleanup EXIT INT TERM
            start_server
            start_worker
            echo ""
            log "═══════════════════════════════════════════════"
            log "  DocEngine is running"
            log "  API Server : http://localhost:7989"
            log "  Swagger UI : http://localhost:7989/docs"
            log "  Celery     : Worker active (4 processes)"
            log "═══════════════════════════════════════════════"
            log "Press Ctrl+C to stop all services."
            echo ""
            wait
            ;;
        *)
            echo "Usage: $0 {server|worker|all|setup|migrate}"
            echo ""
            echo "  server   Start the FastAPI server (default)"
            echo "  worker   Start the Celery worker"
            echo "  all      Start both server and worker"
            echo "  setup    Install dependencies and run migrations"
            echo "  migrate  Run database migrations only"
            exit 1
            ;;
    esac
}

main "$@"