#!/bin/bash
# Automated Rollback Script for Foxhunt HFT Trading System
# Intelligent rollback with automatic trigger detection and recovery validation
#
# This script provides automated rollback capabilities with configurable triggers
# and comprehensive validation of the rollback process.

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
FOXHUNT_HOME="/opt/foxhunt"
RELEASES_DIR="/opt/foxhunt/releases"
CURRENT_LINK="/opt/foxhunt/current"
BACKUP_DIR="/opt/foxhunt/backups"
ROLLBACK_LOG="/home/jgrusewski/Work/foxhunt/logs/rollback-$(date +%s).log"
DEPLOYMENT_MARKER="/tmp/deployment-in-progress"

# Rollback triggers (thresholds)
MAX_LATENCY_US=50
MIN_THROUGHPUT_OPS=500
MAX_ERROR_RATE_PERCENT=1
MAX_CPU_PERCENT=90
MAX_MEMORY_PERCENT=85

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

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

# Create log directory if it doesn't exist
mkdir -p "$(dirname "$ROLLBACK_LOG")"

log() {
    echo -e "${BLUE}[$(date +'%Y-%m-%d %H:%M:%S')]${NC} $1" | tee -a "$ROLLBACK_LOG"
}

error() {
    echo -e "${RED}[ERROR]${NC} $1" | tee -a "$ROLLBACK_LOG"
}

success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1" | tee -a "$ROLLBACK_LOG"
}

warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1" | tee -a "$ROLLBACK_LOG"
}

# Emergency cleanup on script exit
cleanup() {
    local exit_code=$?
    log "Cleaning up rollback artifacts..."
    rm -f "$DEPLOYMENT_MARKER" &>/dev/null || true
    exit $exit_code
}

trap cleanup EXIT INT TERM

usage() {
    echo "Usage: $0 [options]"
    echo "Options:"
    echo "  --trigger <manual|latency|throughput|errors|resources>  Rollback trigger"
    echo "  --target-version <version>                             Specific version to rollback to"
    echo "  --validate-only                                        Only validate rollback capability"
    echo "  --force                                                Force rollback even with warnings"
    echo "  --help                                                 Show this help"
    exit 1
}

# =============================================================================
# TRIGGER DETECTION FUNCTIONS
# =============================================================================

check_latency_trigger() {
    local current_latency=0

    log "Checking latency trigger..."

    # Get current latency from metrics endpoint
    if latency_response=$(curl -s --max-time 5 "http://localhost:8080/metrics" 2>/dev/null); then
        current_latency=$(echo "$latency_response" | grep -o 'foxhunt_order_latency_microseconds [0-9]*' | awk '{print $2}' || echo "0")

        if [ "$current_latency" -gt "$MAX_LATENCY_US" ]; then
            error "Latency trigger activated: ${current_latency}μs > ${MAX_LATENCY_US}μs threshold"
            return 0
        else
            log "Latency acceptable: ${current_latency}μs"
            return 1
        fi
    else
        warning "Could not retrieve latency metrics"
        return 1
    fi
}

check_throughput_trigger() {
    local current_throughput=0

    log "Checking throughput trigger..."

    # Get current throughput from metrics
    if throughput_response=$(curl -s --max-time 5 "http://localhost:8080/metrics" 2>/dev/null); then
        current_throughput=$(echo "$throughput_response" | grep -o 'foxhunt_orders_processed_total [0-9]*' | awk '{print $2}' || echo "0")

        if [ "$current_throughput" -lt "$MIN_THROUGHPUT_OPS" ]; then
            error "Throughput trigger activated: ${current_throughput} ops/min < ${MIN_THROUGHPUT_OPS} threshold"
            return 0
        else
            log "Throughput acceptable: ${current_throughput} ops/min"
            return 1
        fi
    else
        warning "Could not retrieve throughput metrics"
        return 1
    fi
}

check_error_rate_trigger() {
    local error_rate=0

    log "Checking error rate trigger..."

    # Get error rate from metrics
    if error_response=$(curl -s --max-time 5 "http://localhost:8080/metrics" 2>/dev/null); then
        errors=$(echo "$error_response" | grep -o 'foxhunt_errors_total [0-9]*' | awk '{print $2}' || echo "0")
        requests=$(echo "$error_response" | grep -o 'foxhunt_requests_total [0-9]*' | awk '{print $2}' || echo "1")

        if [ "$requests" -gt 0 ]; then
            error_rate=$(echo "scale=2; $errors * 100 / $requests" | bc -l 2>/dev/null || echo "0")

            if (( $(echo "$error_rate > $MAX_ERROR_RATE_PERCENT" | bc -l) )); then
                error "Error rate trigger activated: ${error_rate}% > ${MAX_ERROR_RATE_PERCENT}% threshold"
                return 0
            else
                log "Error rate acceptable: ${error_rate}%"
                return 1
            fi
        fi
    else
        warning "Could not retrieve error rate metrics"
        return 1
    fi
}

check_resource_trigger() {
    log "Checking resource utilization trigger..."

    # Check CPU usage
    local cpu_usage
    cpu_usage=$(top -bn1 | grep "Cpu(s)" | awk '{print $2}' | sed 's/%us,//' | cut -d'%' -f1)

    if (( $(echo "$cpu_usage > $MAX_CPU_PERCENT" | bc -l) )); then
        error "CPU trigger activated: ${cpu_usage}% > ${MAX_CPU_PERCENT}% threshold"
        return 0
    fi

    # Check memory usage
    local memory_usage
    memory_usage=$(free | grep Mem | awk '{printf "%.1f", $3/$2 * 100.0}')

    if (( $(echo "$memory_usage > $MAX_MEMORY_PERCENT" | bc -l) )); then
        error "Memory trigger activated: ${memory_usage}% > ${MAX_MEMORY_PERCENT}% threshold"
        return 0
    fi

    log "Resource utilization acceptable: CPU ${cpu_usage}%, Memory ${memory_usage}%"
    return 1
}

# =============================================================================
# ROLLBACK VALIDATION FUNCTIONS
# =============================================================================

validate_rollback_capability() {
    log "Validating rollback capability..."

    # Check if previous version exists
    if [ ! -f "$FOXHUNT_HOME/.last-good-version" ]; then
        error "No previous version information available for rollback"
        return 1
    fi

    local last_good_version
    last_good_version=$(cat "$FOXHUNT_HOME/.last-good-version")

    local rollback_dir="$RELEASES_DIR/$last_good_version"
    if [ ! -d "$rollback_dir" ]; then
        error "Rollback version directory not found: $rollback_dir"
        return 1
    fi

    # Validate rollback version binaries
    local required_binaries=("foxhunt-core" "foxhunt-tli" "foxhunt-risk" "foxhunt-ml" "foxhunt-data")
    for binary in "${required_binaries[@]}"; do
        if [ ! -x "$rollback_dir/bin/$binary" ]; then
            error "Rollback binary missing or not executable: $binary"
            return 1
        fi
    done

    success "Rollback capability validated - can rollback to version: $last_good_version"
    return 0
}

create_emergency_backup() {
    log "Creating emergency backup before rollback..."

    local backup_timestamp=$(date +%Y%m%d_%H%M%S)
    local emergency_backup="$BACKUP_DIR/emergency_backup_$backup_timestamp"

    mkdir -p "$emergency_backup"

    # Backup current configuration
    if [ -d "$FOXHUNT_HOME/config" ]; then
        cp -r "$FOXHUNT_HOME/config" "$emergency_backup/"
        success "Configuration backed up"
    fi

    # Backup current data state
    if [ -d "$FOXHUNT_HOME/data" ]; then
        cp -r "$FOXHUNT_HOME/data" "$emergency_backup/"
        success "Data state backed up"
    fi

    # Save current system state
    {
        echo "# Emergency backup created: $(date)"
        echo "# Triggered by: $ROLLBACK_TRIGGER"
        echo "# Current version: $(readlink "$CURRENT_LINK" | xargs basename 2>/dev/null || echo "unknown")"
        echo "# System state at backup:"
        systemctl status "${SERVICES[@]}" --no-pager || true
    } > "$emergency_backup/system_state.txt"

    success "Emergency backup created: $emergency_backup"
}

# =============================================================================
# ROLLBACK EXECUTION FUNCTIONS
# =============================================================================

execute_rollback() {
    local target_version="$1"
    local rollback_dir="$RELEASES_DIR/$target_version"

    log "Executing rollback to version: $target_version"

    # Create deployment marker to prevent conflicts
    touch "$DEPLOYMENT_MARKER"

    # Stop services in shutdown order
    log "Stopping services for rollback..."
    for service in "${SHUTDOWN_ORDER[@]}"; do
        log "Stopping $service..."
        if systemctl stop "$service" --timeout=30; then
            success "Stopped $service"
        else
            warning "Failed to stop $service gracefully, forcing stop..."
            systemctl kill "$service" --signal=SIGKILL || true
            sleep 2
        fi
    done

    # Wait for services to fully stop
    sleep 5

    # Update symlink to rollback version
    log "Updating current version symlink..."
    ln -sfn "$rollback_dir" "$CURRENT_LINK"
    success "Updated symlink to rollback version"

    # Start services in dependency order
    log "Starting services with rollback version..."
    for service in "${SERVICES[@]}"; do
        log "Starting $service..."
        if systemctl start "$service"; then
            success "Started $service"
            sleep 3  # Brief pause between service starts
        else
            error "Failed to start $service"
            return 1
        fi
    done

    # Wait for services to initialize
    log "Waiting for services to initialize..."
    sleep 10

    return 0
}

validate_rollback_success() {
    log "Validating rollback success..."

    local validation_errors=0

    # Check service health
    for service in "${SERVICES[@]}"; do
        if systemctl is-active "$service" >/dev/null 2>&1; then
            success "$service is running"
        else
            error "$service is not running after rollback"
            ((validation_errors++))
        fi
    done

    # Check core service health endpoint
    if curl -f -s "http://localhost:8080/health" >/dev/null 2>&1; then
        success "Core service health check passed"
    else
        error "Core service health check failed"
        ((validation_errors++))
    fi

    # Check gRPC connectivity
    if command -v grpcurl >/dev/null 2>&1; then
        if grpcurl -plaintext localhost:50051 list >/dev/null 2>&1; then
            success "gRPC connectivity confirmed"
        else
            error "gRPC connectivity failed"
            ((validation_errors++))
        fi
    fi

    # Basic performance validation
    sleep 5  # Allow metrics to accumulate

    if ! check_latency_trigger && ! check_throughput_trigger; then
        success "Performance metrics acceptable after rollback"
    else
        warning "Performance metrics still degraded after rollback"
        ((validation_errors++))
    fi

    if [ $validation_errors -eq 0 ]; then
        success "Rollback validation passed completely"
        return 0
    else
        error "Rollback validation failed with $validation_errors errors"
        return 1
    fi
}

# =============================================================================
# MAIN EXECUTION LOGIC
# =============================================================================

main() {
    local rollback_trigger="manual"
    local target_version=""
    local validate_only=false
    local force_rollback=false

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --trigger)
                rollback_trigger="$2"
                shift 2
                ;;
            --target-version)
                target_version="$2"
                shift 2
                ;;
            --validate-only)
                validate_only=true
                shift
                ;;
            --force)
                force_rollback=true
                shift
                ;;
            -h|--help)
                usage
                ;;
            *)
                error "Unknown option: $1"
                usage
                ;;
        esac
    done

    # Print banner
    echo "============================================================"
    echo "  Foxhunt HFT Automated Rollback System"
    echo "============================================================"
    echo

    log "Starting automated rollback procedure..."
    log "Trigger: $rollback_trigger"
    log "Validate only: $validate_only"
    log "Force rollback: $force_rollback"

    # Validate rollback capability first
    if ! validate_rollback_capability; then
        error "Rollback capability validation failed"
        exit 1
    fi

    # If validate-only mode, exit here
    if [ "$validate_only" = true ]; then
        success "Rollback capability validation completed successfully"
        exit 0
    fi

    # Determine target version
    if [ -z "$target_version" ]; then
        if [ -f "$FOXHUNT_HOME/.last-good-version" ]; then
            target_version=$(cat "$FOXHUNT_HOME/.last-good-version")
            log "Using last known good version: $target_version"
        else
            error "No target version specified and no last good version available"
            exit 1
        fi
    fi

    # Check rollback triggers (unless forced or manual)
    local should_rollback=false

    if [ "$rollback_trigger" = "manual" ] || [ "$force_rollback" = true ]; then
        should_rollback=true
        log "Manual rollback or force flag - proceeding without trigger validation"
    else
        case $rollback_trigger in
            latency)
                if check_latency_trigger; then
                    should_rollback=true
                fi
                ;;
            throughput)
                if check_throughput_trigger; then
                    should_rollback=true
                fi
                ;;
            errors)
                if check_error_rate_trigger; then
                    should_rollback=true
                fi
                ;;
            resources)
                if check_resource_trigger; then
                    should_rollback=true
                fi
                ;;
            *)
                # Check all triggers
                if check_latency_trigger || check_throughput_trigger || check_error_rate_trigger || check_resource_trigger; then
                    should_rollback=true
                fi
                ;;
        esac
    fi

    if [ "$should_rollback" = false ]; then
        success "No rollback triggers activated - system appears healthy"
        exit 0
    fi

    # Create emergency backup
    create_emergency_backup

    # Execute rollback
    log "Initiating rollback to version: $target_version"

    if execute_rollback "$target_version"; then
        success "Rollback execution completed"
    else
        error "Rollback execution failed"
        exit 1
    fi

    # Validate rollback success
    if validate_rollback_success; then
        success "Rollback completed successfully and validated"

        # Update last good version to the rollback version
        echo "$target_version" > "$FOXHUNT_HOME/.last-good-version"

        # Clean up deployment marker
        rm -f "$DEPLOYMENT_MARKER"

        log "Rollback procedure completed successfully"
        log "Rollback log: $ROLLBACK_LOG"

        exit 0
    else
        error "Rollback validation failed - system may require manual intervention"
        log "Rollback log: $ROLLBACK_LOG"
        exit 1
    fi
}

# Set global variable for cleanup
ROLLBACK_TRIGGER="$1"

# Execute main function
main "$@"