#!/bin/bash
set -e

# Hero System Startup Script - Wrapper
# Runs individual run.sh scripts in background with logging

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

# Base directory
HERO_BASE="/Users/timurgordon/code/git.ourworld.tf/herocode"

# Log directory
LOG_DIR="$HERO_BASE/home/logs"
mkdir -p "$LOG_DIR"

# PID file directory
PID_DIR="$HERO_BASE/home/pids"
mkdir -p "$PID_DIR"

# Function to print colored messages
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

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

# Function to check if Redis is running
check_redis() {
    log_info "Checking Redis connection..."
    if redis-cli ping > /dev/null 2>&1; then
        log_success "Redis is running"
        return 0
    else
        log_error "Redis is not running. Please start Redis first."
        log_info "You can start Redis with: redis-server"
        return 1
    fi
}

# Function to start a component with log prefix
start_component() {
    local name=$1
    local script=$2
    local log_file=$3
    local pid_file=$4
    local prefix=$5
    
    log_info "Starting $name..."
    
    # Start the script and pipe output through a prefix filter
    "$script" 2>&1 | while IFS= read -r line; do
        echo -e "${prefix}${NC} $line"
    done &
    
    local pid=$!
    echo $pid > "$pid_file"
    sleep 2
    
    if ps -p $pid > /dev/null 2>&1; then
        log_success "$name started (PID: $pid)"
    else
        log_error "Failed to start $name"
        return 1
    fi
}

# Function to stop a process
stop_process() {
    local name=$1
    local pid_file=$2
    
    if [ -f "$pid_file" ]; then
        local pid=$(cat "$pid_file")
        if ps -p $pid > /dev/null 2>&1; then
            log_info "Stopping $name (PID: $pid)..."
            kill "$pid" 2>/dev/null || true
            sleep 2
            
            # Force kill if still running
            if ps -p "$pid" > /dev/null 2>&1; then
                log_warning "Force killing $name..."
                kill -9 "$pid" 2>/dev/null || true
            fi
            
            rm -f "$pid_file"
            log_success "$name stopped"
        else
            log_info "$name is not running"
            rm -f "$pid_file"
        fi
    else
        log_info "$name is not running"
    fi
}

# Function to show status
show_status() {
    echo ""
    log_info "Hero System Status:"
    echo ""
    
    for service in coordinator supervisor osiris_runner; do
        local pid_file="$PID_DIR/${service}.pid"
        if [ -f "$pid_file" ]; then
            local pid=$(cat "$pid_file")
            if ps -p $pid > /dev/null 2>&1; then
                echo -e "  ${GREEN}●${NC} $service (PID: $pid)"
            else
                echo -e "  ${RED}●${NC} $service (dead, PID file exists)"
            fi
        else
            echo -e "  ${RED}●${NC} $service (not running)"
        fi
    done
    echo ""
}

# Function to stop all services
stop_all() {
    log_info "Stopping Hero System..."
    stop_process "Osiris Runner" "$PID_DIR/osiris_runner.pid"
    stop_process "Supervisor" "$PID_DIR/supervisor.pid"
    stop_process "Coordinator" "$PID_DIR/coordinator.pid"
    log_success "All services stopped"
}

# Trap handler for cleanup on exit
cleanup() {
    echo ""
    log_warning "Received exit signal, stopping all services..."
    stop_all
    exit 0
}

# Main script
case "${1:-start}" in
    start)
        log_info "Starting Hero System..."
        
        # Set up trap for SIGINT and SIGTERM
        trap cleanup SIGINT SIGTERM
        
        # Check Redis
        if ! check_redis; then
            exit 1
        fi
        
        # # Start services
        # start_component "Coordinator" \
        #     "$HERO_BASE/herocoordinator/scripts/run.sh" \
        #     "$LOG_DIR/coordinator.log" \
        #     "$PID_DIR/coordinator.pid" \
        #     "${GREEN}[COORD]" || exit 1
            
        start_component "Supervisor" \
            "$HERO_BASE/supervisor/scripts/run.sh" \
            "$LOG_DIR/supervisor.log" \
            "$PID_DIR/supervisor.pid" \
            "${BLUE}[SUPER]" || exit 1
            
        start_component "Osiris Runner" \
            "$HERO_BASE/runner_rust/scripts/run.sh" \
            "$LOG_DIR/osiris_runner.log" \
            "$PID_DIR/osiris_runner.pid" \
            "${YELLOW}[OSIRS]" || exit 1
        
        echo ""
        log_success "Hero System started successfully!"
        show_status
        
        echo ""
        log_info "Displaying live logs from all services..."
        log_info "Press Ctrl+C to stop all services"
        echo ""
        
        # Keep running and wait for signals
        wait
        ;;
        
    stop)
        stop_all
        ;;
        
    restart)
        stop_all
        sleep 2
        $0 start
        ;;
        
    status)
        show_status
        ;;
        
    *)
        echo "Usage: $0 {start|stop|restart|status}"
        echo ""
        echo "Commands:"
        echo "  start   - Start all Hero services with live log display (default)"
        echo "  stop    - Stop all Hero services"
        echo "  restart - Restart all Hero services"
        echo "  status  - Show status of all services"
        echo ""
        echo "Note: 'start' displays live logs from all services with color-coded prefixes."
        echo "      Each service's run.sh will build itself before starting."
        exit 1
        ;;
esac