#!/bin/bash
# Emergency rollback script for Foxhunt HFT Trading System
# Implements rapid rollback with minimal downtime for critical production issues

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
EMERGENCY_LOG="/home/jgrusewski/Work/foxhunt/logs/emergency-rollback-$(date +%s).log"
FOXHUNT_HOME="/opt/foxhunt"
RELEASES_DIR="/opt/foxhunt/releases"
CURRENT_LINK="/opt/foxhunt/current"
ALERT_WEBHOOK="${FOXHUNT_ALERT_WEBHOOK:-}"

# Services in dependency order (reverse for shutdown)
SERVICES=("foxhunt-tli" "foxhunt-ml" "foxhunt-risk" "foxhunt-data" "foxhunt-core")
SHUTDOWN_SERVICES=("foxhunt-core" "foxhunt-data" "foxhunt-risk" "foxhunt-ml" "foxhunt-tli")

# Emergency thresholds
MAX_ROLLBACK_TIME_SECONDS=60
HEALTH_CHECK_TIMEOUT=10

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

log() {
    local level="${2:-INFO}"
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [$level] $1" | tee -a "$EMERGENCY_LOG"
}

error() {
    log "$1" "ERROR"
    echo -e "${RED}EMERGENCY: $1${NC}" >&2
}

success() {
    log "$1" "SUCCESS"
    echo -e "${GREEN}SUCCESS: $1${NC}"
}

warning() {
    log "$1" "WARNING"
    echo -e "${YELLOW}WARNING: $1${NC}"
}

critical() {
    log "$1" "CRITICAL"
    echo -e "${BOLD}${RED}CRITICAL: $1${NC}" >&2
    send_alert "CRITICAL: $1"
}

info() {
    log "$1" "INFO"
    echo -e "${BLUE}INFO: $1${NC}"
}

send_alert() {
    local message="$1"
    
    if [ -n "$ALERT_WEBHOOK" ]; then
        curl -s -X POST "$ALERT_WEBHOOK" \
             -H "Content-Type: application/json" \
             -d "{\"text\":\"🚨 Foxhunt Emergency Rollback: $message\"}" \
             > /dev/null 2>&1 || true
    fi
    
    # Also send to system log
    logger -p daemon.crit "Foxhunt Emergency Rollback: $message"
}

cleanup() {
    log "Emergency rollback script completed with exit code ${1:-1}"
    
    if [ "${1:-1}" -eq 0 ]; then
        send_alert "Emergency rollback completed successfully"
    else
        send_alert "Emergency rollback failed - manual intervention required"
    fi
}

trap cleanup EXIT INT TERM

usage() {
    echo "Usage: $0 [options]"
    echo ""
    echo "Options:"
    echo "  --reason REASON      Reason for emergency rollback (required)"
    echo "  --validate-only      Only validate rollback capability, don't execute"
    echo "  --skip-health        Skip health checks (use only in extreme emergencies)"
    echo "  --force              Force rollback even if risks are detected"
    echo "  -h, --help          Show this help message"
    echo ""
    echo "This script performs an emergency rollback to the last known good version."
    echo "It prioritizes speed over safety checks - use only in critical situations."
    exit 1
}

detect_current_environment() {
    log "Detecting current deployment environment..."
    
    local current_env="unknown"
    
    # Check if blue-green deployment is active
    if [ -f "$FOXHUNT_HOME/.active-environment" ]; then
        current_env=$(cat "$FOXHUNT_HOME/.active-environment")
    elif [ -L "$CURRENT_LINK" ]; then
        local link_target
        link_target=$(readlink "$CURRENT_LINK")
        if [[ "$link_target" == *"blue"* ]]; then
            current_env="blue"
        elif [[ "$link_target" == *"green"* ]]; then
            current_env="green"
        fi
    fi
    
    info "Current environment: $current_env"
    echo "$current_env"
}

get_last_good_version() {
    log "Identifying last known good version..."
    
    local last_good=""
    
    # Check for explicit last good version marker
    if [ -f "$FOXHUNT_HOME/.last-good-version" ]; then
        last_good=$(cat "$FOXHUNT_HOME/.last-good-version")
    fi
    
    # Fallback: get previous version from releases directory
    if [ -z "$last_good" ] && [ -d "$RELEASES_DIR" ]; then
        # Get the second most recent release
        last_good=$(ls -t "$RELEASES_DIR" | head -2 | tail -1)
    fi
    
    if [ -z "$last_good" ]; then
        critical "Cannot determine last good version for rollback"
        return 1
    fi
    
    local rollback_dir="$RELEASES_DIR/$last_good"
    if [ ! -d "$rollback_dir" ]; then
        critical "Last good version directory not found: $rollback_dir"
        return 1
    fi
    
    info "Last known good version: $last_good"
    echo "$last_good"
}

rapid_health_check() {
    local service_name="$1"
    local health_url="$2"
    local timeout="${3:-$HEALTH_CHECK_TIMEOUT}"
    
    # Rapid health check with minimal retries
    local attempt=0
    local max_attempts=3
    
    while [ $attempt -lt $max_attempts ]; do
        if timeout "$timeout" curl -f -s "$health_url" > /dev/null 2>&1; then
            return 0
        fi
        
        attempt=$((attempt + 1))
        sleep 1
    done
    
    return 1
}

emergency_stop_services() {
    log "Emergency stopping all services..."
    
    local stop_start=$(date +%s)
    
    # Try graceful shutdown first (parallel)
    for service in "${SHUTDOWN_SERVICES[@]}"; do
        (
            log "Gracefully stopping $service..."
            systemctl stop "$service" || true
        ) &
    done
    
    # Wait for graceful shutdown (max 15 seconds)
    local graceful_timeout=15
    local elapsed=0
    
    while [ $elapsed -lt $graceful_timeout ]; do
        local all_stopped=true
        
        for service in "${SHUTDOWN_SERVICES[@]}"; do
            if systemctl is-active "$service" > /dev/null 2>&1; then
                all_stopped=false
                break
            fi
        done
        
        if [ "$all_stopped" = true ]; then
            break
        fi
        
        sleep 1
        elapsed=$((elapsed + 1))
    done
    
    # Force kill any remaining processes
    for service in "${SHUTDOWN_SERVICES[@]}"; do
        if systemctl is-active "$service" > /dev/null 2>&1; then
            warning "Force killing $service..."
            systemctl kill -s KILL "$service" || true
        fi
    done
    
    # Stop Docker containers if they exist
    if command -v docker > /dev/null 2>&1; then
        log "Stopping Docker containers..."
        docker stop $(docker ps -q --filter "name=foxhunt") 2>/dev/null || true
    fi
    
    local stop_end=$(date +%s)
    local stop_duration=$((stop_end - stop_start))
    
    info "Services stopped in ${stop_duration}s"
}

emergency_rollback_execution() {
    local rollback_version="$1"
    local skip_health="$2"
    local rollback_dir="$RELEASES_DIR/$rollback_version"
    
    log "Executing emergency rollback to version $rollback_version..."
    
    local rollback_start=$(date +%s)
    
    # Update symlink atomically
    local temp_link="${CURRENT_LINK}.emergency.$$"
    ln -s "$rollback_dir" "$temp_link"
    mv "$temp_link" "$CURRENT_LINK"
    
    # Start services in dependency order (parallel where safe)
    log "Starting services with rollback version..."
    
    # Start core infrastructure services first
    for service in "foxhunt-core" "foxhunt-data"; do
        log "Starting $service..."
        systemctl start "$service"
        
        # Quick health check for critical services
        case $service in
            foxhunt-core)
                if [ "$skip_health" != "true" ]; then
                    if ! rapid_health_check "$service" "http://localhost:8080/health" 5; then
                        critical "$service failed to start properly after rollback"
                        return 1
                    fi
                fi
                ;;
        esac
        
        sleep 2  # Brief pause between critical services
    done
    
    # Start remaining services in parallel
    for service in "foxhunt-risk" "foxhunt-ml" "foxhunt-tli"; do
        (
            log "Starting $service..."
            systemctl start "$service"
        ) &
    done
    
    # Wait for all background starts
    wait
    
    # Final validation
    sleep 5
    
    if [ "$skip_health" != "true" ]; then
        log "Performing rapid system validation..."
        
        local validation_failed=false
        
        # Check critical services
        if ! rapid_health_check "foxhunt-core" "http://localhost:8080/health" 3; then
            error "Core service health check failed after rollback"
            validation_failed=true
        fi
        
        if ! rapid_health_check "foxhunt-tli" "http://localhost:8081/health" 3; then
            warning "TLI service health check failed after rollback"
        fi
        
        # Check gRPC connectivity
        if ! timeout 3 grpcurl -plaintext localhost:50051 list > /dev/null 2>&1; then
            warning "gRPC connectivity check failed after rollback"
        fi
        
        if [ "$validation_failed" = true ]; then
            critical "Critical service validation failed after rollback"
            return 1
        fi
    fi
    
    local rollback_end=$(date +%s)
    local rollback_duration=$((rollback_end - rollback_start))
    
    # Update version markers
    echo "$rollback_version" > "$FOXHUNT_HOME/.current-version"
    
    success "Emergency rollback completed in ${rollback_duration}s"
    
    # Check if we met our time target
    if [ $rollback_duration -le $MAX_ROLLBACK_TIME_SECONDS ]; then
        success "Rollback completed within target time (${MAX_ROLLBACK_TIME_SECONDS}s)"
    else
        warning "Rollback took longer than target (${rollback_duration}s > ${MAX_ROLLBACK_TIME_SECONDS}s)"
    fi
    
    return 0
}

validate_rollback_capability() {
    log "Validating emergency rollback capability..."
    
    local validation_errors=0
    
    # Check if we can determine last good version
    local last_good
    if ! last_good=$(get_last_good_version); then
        error "Cannot determine last good version"
        validation_errors=$((validation_errors + 1))
    fi
    
    # Check releases directory
    if [ ! -d "$RELEASES_DIR" ]; then
        error "Releases directory not found: $RELEASES_DIR"
        validation_errors=$((validation_errors + 1))
    fi
    
    # Check systemctl availability
    if ! command -v systemctl > /dev/null 2>&1; then
        error "systemctl command not available"
        validation_errors=$((validation_errors + 1))
    fi
    
    # Check service definitions
    for service in "${SERVICES[@]}"; do
        if ! systemctl is-enabled "$service" > /dev/null 2>&1; then
            warning "Service $service is not enabled"
        fi
    done
    
    # Check disk space
    local available_space
    available_space=$(df "$FOXHUNT_HOME" | awk 'NR==2 {print $4}')
    if [ "$available_space" -lt 1048576 ]; then  # Less than 1GB
        error "Insufficient disk space for rollback operations"
        validation_errors=$((validation_errors + 1))
    fi
    
    # Check permissions
    if [ ! -w "$FOXHUNT_HOME" ]; then
        error "No write permissions to Foxhunt home directory"
        validation_errors=$((validation_errors + 1))
    fi
    
    if [ $validation_errors -eq 0 ]; then
        success "Emergency rollback capability validation passed"
        return 0
    else
        error "Emergency rollback capability validation failed ($validation_errors errors)"
        return 1
    fi
}

main() {
    local reason=""
    local validate_only=false
    local skip_health=false
    local force=false
    
    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --reason)
                reason="$2"
                shift 2
                ;;
            --validate-only)
                validate_only=true
                shift
                ;;
            --skip-health)
                skip_health=true
                shift
                ;;
            --force)
                force=true
                shift
                ;;
            -h|--help)
                usage
                ;;
            *)
                error "Unknown option: $1"
                usage
                ;;
        esac
    done
    
    # Validate reason is provided (unless validate-only)
    if [ "$validate_only" = false ] && [ -z "$reason" ]; then
        error "Emergency rollback reason is required"
        usage
    fi
    
    critical "EMERGENCY ROLLBACK INITIATED"
    log "Script started by: $(whoami)"
    log "Reason: ${reason:-validation-only}"
    log "Validate only: $validate_only"
    log "Skip health: $skip_health"
    log "Force: $force"
    
    # Always validate capability first
    if ! validate_rollback_capability; then
        if [ "$force" != "true" ]; then
            critical "Rollback capability validation failed - use --force to override"
            return 1
        else
            warning "Proceeding with rollback despite validation failures"
        fi
    fi
    
    # If validation-only mode, exit here
    if [ "$validate_only" = true ]; then
        success "Emergency rollback capability validated"
        return 0
    fi
    
    # Get rollback target
    local rollback_version
    if ! rollback_version=$(get_last_good_version); then
        critical "Cannot proceed with rollback - no valid target version"
        return 1
    fi
    
    # Record rollback initiation
    send_alert "Emergency rollback initiated - Reason: $reason - Target: $rollback_version"
    
    # Emergency stop all services
    emergency_stop_services
    
    # Execute rollback
    if emergency_rollback_execution "$rollback_version" "$skip_health"; then
        success "Emergency rollback to version $rollback_version completed successfully"
        send_alert "Emergency rollback completed successfully - Version: $rollback_version"
        return 0
    else
        critical "Emergency rollback failed - manual intervention required"
        return 1
    fi
}

# Execute main function
main "$@"