#!/usr/bin/env bash
#===============================================================================
#
#   WISDOM - AI Services Management Script
#
#===============================================================================
#
#   DESCRIPTION:
#       This script manages the AI backend services used by the "wisdom" system.
#       It provides commands to start, stop, restart, and check the status of:
#
#       1. OLLAMA SERVICE (Port 11434)
#          - Local LLM inference server
#          - Runs models like deepseek-r1:14b, qwen2.5:14b
#          - Managed via systemd (ollama.service)
#          - API endpoint: http://localhost:11434
#
#       2. FASTER-WHISPER SERVER (Port 8000)
#          - Speech-to-text transcription server
#          - Uses OpenAI Whisper models with CTranslate2 optimization
#          - Runs as a uvicorn ASGI server
#          - Location: /root/faster-whisper-server
#          - API endpoint: http://localhost:8000
#
#   USAGE:
#       wisdom start    | -l   - Start/launch all AI services
#       wisdom stop     | -c   - Stop/close all AI services (frees RAM)
#       wisdom restart  | -r   - Restart all AI services
#       wisdom status   | -s   - Show status of all services
#       wisdom help     | -h   - Show this help message
#
#   REQUIREMENTS:
#       - sudo access (for whisper server management)
#       - systemd (for ollama service)
#       - /root/faster-whisper-server directory with venv
#
#   PORTS USED:
#       - 11434: Ollama API (localhost only)
#       - 8000:  Faster-Whisper API
#
#   MEMORY USAGE:
#       - Ollama: ~200MB idle, up to 8GB+ when running large models
#       - Whisper: ~500MB-2GB depending on model loaded
#
#   AUTHOR: Generated for krisyotam's wisdom system
#   DATE:   2026-01-19
#
#===============================================================================

set -e  # Exit on error

#-------------------------------------------------------------------------------
# CONFIGURATION
#-------------------------------------------------------------------------------

# Ollama settings
OLLAMA_SERVICE="ollama.service"
OLLAMA_PORT=11434
OLLAMA_HOST="127.0.0.1"

# Faster-Whisper settings
WHISPER_DIR="/root/faster-whisper-server"
WHISPER_VENV="$WHISPER_DIR/.venv"
WHISPER_PORT=8000

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

#-------------------------------------------------------------------------------
# HELPER FUNCTIONS
#-------------------------------------------------------------------------------

# Print colored status messages
print_status() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[OK]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Check if a service is running on a specific port
check_port() {
    local port=$1
    ss -tlnp 2>/dev/null | grep -q ":$port " && return 0 || return 1
}

# Get PID of process on port
get_pid_on_port() {
    local port=$1
    ss -tlnp 2>/dev/null | grep ":$port " | grep -oP 'pid=\K\d+' | head -1
}

#-------------------------------------------------------------------------------
# SERVICE MANAGEMENT FUNCTIONS
#-------------------------------------------------------------------------------

#
# START OLLAMA
# Uses systemd to start the ollama service
#
start_ollama() {
    print_status "Starting Ollama service..."

    if systemctl is-active --quiet "$OLLAMA_SERVICE"; then
        print_warning "Ollama is already running"
    else
        sudo systemctl start "$OLLAMA_SERVICE"
        sleep 2

        if systemctl is-active --quiet "$OLLAMA_SERVICE"; then
            print_success "Ollama started successfully on port $OLLAMA_PORT"
        else
            print_error "Failed to start Ollama"
            return 1
        fi
    fi
}

#
# STOP OLLAMA
# Uses systemd to stop the ollama service
#
stop_ollama() {
    print_status "Stopping Ollama service..."

    if ! systemctl is-active --quiet "$OLLAMA_SERVICE"; then
        print_warning "Ollama is not running"
    else
        sudo systemctl stop "$OLLAMA_SERVICE"
        print_success "Ollama stopped"
    fi
}

#
# START FASTER-WHISPER SERVER
# Starts the uvicorn server in the background
# Requires sudo since it runs from /root
#
start_whisper() {
    print_status "Starting Faster-Whisper server..."

    # Check if already running
    if check_port $WHISPER_PORT; then
        print_warning "Whisper server is already running on port $WHISPER_PORT"
        return 0
    fi

    # Check if the directory exists
    if [ ! -d "$WHISPER_DIR" ]; then
        print_error "Whisper server directory not found: $WHISPER_DIR"
        return 1
    fi

    # Start the server using sudo
    # The server runs in the background with nohup
    sudo bash -c "
        cd $WHISPER_DIR && \
        source $WHISPER_VENV/bin/activate && \
        nohup uv run uvicorn --factory faster_whisper_server.main:create_app \
            --host 0.0.0.0 \
            --port $WHISPER_PORT \
            > /var/log/whisper-server.log 2>&1 &
    "

    # Wait for server to start
    sleep 3

    if check_port $WHISPER_PORT; then
        print_success "Whisper server started on port $WHISPER_PORT"
    else
        print_error "Failed to start Whisper server"
        print_status "Check logs: sudo cat /var/log/whisper-server.log"
        return 1
    fi
}

#
# STOP FASTER-WHISPER SERVER
# Finds and kills the uvicorn process
#
stop_whisper() {
    print_status "Stopping Faster-Whisper server..."

    # Find processes related to faster-whisper-server
    local pids=$(pgrep -f "faster_whisper_server" 2>/dev/null || true)

    if [ -z "$pids" ]; then
        print_warning "Whisper server is not running"
        return 0
    fi

    # Kill the processes
    for pid in $pids; do
        sudo kill "$pid" 2>/dev/null || true
    done

    sleep 1

    # Force kill if still running
    pids=$(pgrep -f "faster_whisper_server" 2>/dev/null || true)
    if [ -n "$pids" ]; then
        for pid in $pids; do
            sudo kill -9 "$pid" 2>/dev/null || true
        done
    fi

    print_success "Whisper server stopped"
}

#
# SHOW STATUS OF ALL SERVICES
#
show_status() {
    echo ""
    echo "========================================"
    echo "       WISDOM AI SERVICES STATUS        "
    echo "========================================"
    echo ""

    # Ollama status
    echo -e "${BLUE}OLLAMA SERVICE${NC}"
    echo "  Port: $OLLAMA_PORT"
    if systemctl is-active --quiet "$OLLAMA_SERVICE"; then
        echo -e "  Status: ${GREEN}RUNNING${NC}"
        # Show memory usage
        local ollama_mem=$(systemctl show "$OLLAMA_SERVICE" --property=MemoryCurrent 2>/dev/null | cut -d= -f2)
        if [ -n "$ollama_mem" ] && [ "$ollama_mem" != "[not set]" ]; then
            echo "  Memory: $(numfmt --to=iec $ollama_mem 2>/dev/null || echo $ollama_mem)"
        fi
        # Show loaded models
        echo "  Models:"
        curl -s "http://$OLLAMA_HOST:$OLLAMA_PORT/api/tags" 2>/dev/null | \
            grep -oP '"name":"[^"]+' | cut -d'"' -f4 | while read model; do
                echo "    - $model"
            done
    else
        echo -e "  Status: ${RED}STOPPED${NC}"
    fi
    echo ""

    # Whisper status
    echo -e "${BLUE}FASTER-WHISPER SERVER${NC}"
    echo "  Port: $WHISPER_PORT"
    if check_port $WHISPER_PORT; then
        echo -e "  Status: ${GREEN}RUNNING${NC}"
        local pid=$(get_pid_on_port $WHISPER_PORT)
        if [ -n "$pid" ]; then
            local mem=$(ps -o rss= -p $pid 2>/dev/null | awk '{print $1*1024}')
            echo "  Memory: $(numfmt --to=iec $mem 2>/dev/null || echo "unknown")"
        fi
    else
        echo -e "  Status: ${RED}STOPPED${NC}"
    fi
    echo ""
    echo "========================================"
}

#
# SHOW HELP MESSAGE
#
show_help() {
    echo ""
    echo "WISDOM - AI Services Management"
    echo ""
    echo "Usage: wisdom <command>"
    echo ""
    echo "Commands:"
    echo "  start,   -l    Start/launch all AI services (Ollama + Whisper)"
    echo "  stop,    -c    Stop/close all AI services (frees RAM)"
    echo "  restart, -r    Restart all AI services"
    echo "  status,  -s    Show status of all services"
    echo "  help,    -h    Show this help message"
    echo ""
    echo "Services managed:"
    echo "  - Ollama (LLM inference) - Port 11434"
    echo "  - Faster-Whisper (Speech-to-text) - Port 8000"
    echo ""
}

#-------------------------------------------------------------------------------
# MAIN COMMAND HANDLER
#-------------------------------------------------------------------------------

case "${1:-help}" in
    start|-l|--launch)
        echo ""
        print_status "Starting WISDOM AI services..."
        echo ""
        start_ollama
        start_whisper
        echo ""
        print_success "All services started!"
        show_status
        ;;

    stop|-c|--close)
        echo ""
        print_status "Stopping WISDOM AI services..."
        echo ""
        stop_whisper
        stop_ollama
        echo ""
        print_success "All services stopped. RAM freed."
        ;;

    restart|-r|--restart)
        echo ""
        print_status "Restarting WISDOM AI services..."
        echo ""
        stop_whisper
        stop_ollama
        sleep 2
        start_ollama
        start_whisper
        echo ""
        print_success "All services restarted!"
        show_status
        ;;

    status|-s|--status)
        show_status
        ;;

    help|--help|-h)
        show_help
        ;;

    *)
        print_error "Unknown command: $1"
        show_help
        exit 1
        ;;
esac