#!/bin/bash
# Staging Deployment Script for Foxhunt HFT Trading System
# Automated staging environment deployment for testing and validation
#
# This script manages staging deployments for testing new versions
# before production rollout.

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
STAGING_LOG="/home/jgrusewski/Work/foxhunt/logs/staging-deployment-$(date +%s).log"
STAGING_HOME="/opt/foxhunt/staging"
COMPOSE_FILE="$SCRIPT_DIR/../docker/docker-compose.staging.yml"
STAGING_NETWORK="foxhunt-staging-net"

# 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
mkdir -p "$(dirname "$STAGING_LOG")"

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

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

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

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

usage() {
    echo "Usage: $0 <command> [options]"
    echo "Commands:"
    echo "  deploy <version>     Deploy version to staging"
    echo "  test                 Run staging tests"
    echo "  status               Show staging environment status"
    echo "  logs <service>       Show logs for service"
    echo "  cleanup              Clean up staging environment"
    echo "  reset                Reset staging environment"
    echo ""
    echo "Options:"
    echo "  --with-load-test     Include load testing"
    echo "  --skip-health        Skip health checks"
    echo "  --force              Force deployment even with warnings"
    echo "  --help               Show this help"
    exit 1
}

# =============================================================================
# STAGING ENVIRONMENT MANAGEMENT
# =============================================================================

setup_staging_environment() {
    log "Setting up staging environment..."

    # Create staging directories
    mkdir -p "$STAGING_HOME"/{config,data,models,checkpoints,logs}

    # Create staging configuration
    cat > "$STAGING_HOME/config/staging.toml" <<EOF
[general]
environment = "staging"
log_level = "debug"
test_mode = true

[database]
url = "postgresql://foxhunt_staging:staging_pass@localhost/foxhunt_staging"
pool_size = 5

[trading]
max_latency_us = 100
test_mode = true
paper_trading = true
max_position_size = 1000

[risk_management]
max_daily_loss = 1000.0
position_limit = 10000.0
test_mode = true

[ml_training]
gpu_enabled = true
batch_size = 16
max_epochs = 100
model_save_interval = 10

[monitoring]
prometheus_endpoint = "http://prometheus-staging:9090"
metrics_interval = 5

[market_data]
polygon_api_key = "${POLYGON_STAGING_API_KEY:-demo}"
paper_trading = true
EOF

    success "Staging environment setup completed"
}

deploy_to_staging() {
    local version="$1"
    local with_load_test=false
    local skip_health=false
    local force_deploy=false

    # Parse additional options
    shift
    while [[ $# -gt 0 ]]; do
        case $1 in
            --with-load-test)
                with_load_test=true
                shift
                ;;
            --skip-health)
                skip_health=true
                shift
                ;;
            --force)
                force_deploy=true
                shift
                ;;
            *)
                error "Unknown option: $1"
                return 1
                ;;
        esac
    done

    log "Deploying version $version to staging environment..."

    # Setup staging environment
    setup_staging_environment

    # Stop existing staging services
    log "Stopping existing staging services..."
    docker-compose -f "$COMPOSE_FILE" down || true

    # Pull/build new images
    log "Building staging images for version $version..."

    # Set version in environment for Docker builds
    export FOXHUNT_VERSION="$version"

    if docker-compose -f "$COMPOSE_FILE" build; then
        success "Staging images built successfully"
    else
        error "Failed to build staging images"
        return 1
    fi

    # Start staging services
    log "Starting staging services..."
    if docker-compose -f "$COMPOSE_FILE" up -d; then
        success "Staging services started"
    else
        error "Failed to start staging services"
        return 1
    fi

    # Wait for services to stabilize
    log "Waiting for services to stabilize..."
    sleep 30

    # Health checks
    if [ "$skip_health" = false ]; then
        log "Running health checks..."
        if run_staging_health_checks; then
            success "All staging health checks passed"
        else
            if [ "$force_deploy" = false ]; then
                error "Health checks failed - deployment aborted"
                return 1
            else
                warning "Health checks failed but continuing due to --force"
            fi
        fi
    fi

    # Load testing
    if [ "$with_load_test" = true ]; then
        log "Starting load testing..."
        docker-compose -f "$COMPOSE_FILE" --profile testing up -d load-generator
        success "Load testing started"
    fi

    success "Staging deployment completed for version $version"
    log "Staging services available at:"
    log "  Core Service: http://localhost:8090"
    log "  TLI Service: http://localhost:8091 (gRPC: localhost:50061)"
    log "  ML Service: http://localhost:8092"
    log "  Risk Service: http://localhost:8093"
    log "  Data Service: http://localhost:8094"
    log "  Prometheus: http://localhost:9191"
    log "  TensorBoard: http://localhost:6016"
}

run_staging_health_checks() {
    local health_endpoints=(
        "http://localhost:8090/health"  # Core
        "http://localhost:8091/health"  # TLI
        "http://localhost:8092/health"  # ML
        "http://localhost:8093/health"  # Risk
        "http://localhost:8094/health"  # Data
    )

    local healthy_count=0
    local total_endpoints=${#health_endpoints[@]}

    for endpoint in "${health_endpoints[@]}"; do
        local service_name=$(echo "$endpoint" | sed 's/.*:\([0-9]*\).*/Port \1/')

        # Wait up to 60 seconds for each service
        local attempts=0
        local max_attempts=12

        while [ $attempts -lt $max_attempts ]; do
            if curl -f -s --max-time 5 "$endpoint" >/dev/null 2>&1; then
                success "Health check passed: $service_name"
                ((healthy_count++))
                break
            fi

            ((attempts++))
            if [ $attempts -lt $max_attempts ]; then
                log "Health check attempt $attempts/$max_attempts failed for $service_name, retrying..."
                sleep 5
            fi
        done

        if [ $attempts -eq $max_attempts ]; then
            error "Health check failed for $service_name after $max_attempts attempts"
        fi
    done

    log "Health check summary: $healthy_count/$total_endpoints services healthy"

    if [ "$healthy_count" -eq "$total_endpoints" ]; then
        return 0
    else
        return 1
    fi
}

run_staging_tests() {
    log "Running staging environment tests..."

    # Basic connectivity tests
    log "Testing service connectivity..."
    if run_staging_health_checks; then
        success "Service connectivity tests passed"
    else
        error "Service connectivity tests failed"
    fi

    # gRPC connectivity test
    log "Testing gRPC connectivity..."
    if command -v grpcurl >/dev/null 2>&1; then
        if grpcurl -plaintext localhost:50061 list >/dev/null 2>&1; then
            success "gRPC connectivity test passed"
        else
            error "gRPC connectivity test failed"
        fi
    else
        warning "grpcurl not available - skipping gRPC test"
    fi

    # Performance test
    log "Running performance test..."
    if [ -x "$SCRIPT_DIR/performance-benchmark.sh" ]; then
        # Run a quick performance test against staging
        local staging_endpoint="http://localhost:8090"

        # Simple latency test
        local start_time end_time latency_ms
        start_time=$(date +%s%N)
        if curl -f -s --max-time 5 "$staging_endpoint/health" >/dev/null 2>&1; then
            end_time=$(date +%s%N)
            latency_ms=$(((end_time - start_time) / 1000000))
            log "Staging service latency: ${latency_ms}ms"

            if [ "$latency_ms" -lt 100 ]; then
                success "Performance test passed"
            else
                warning "High latency detected: ${latency_ms}ms"
            fi
        else
            error "Performance test failed - service not responding"
        fi
    else
        warning "Performance benchmark script not available"
    fi

    # ML service specific tests
    log "Testing ML Training Service integration..."
    if curl -f -s --max-time 10 "http://localhost:8092/health" >/dev/null 2>&1; then
        success "ML Training Service test passed"

        # Check if GPU is available in staging
        if docker exec foxhunt-ml-training-staging nvidia-smi >/dev/null 2>&1; then
            success "GPU access confirmed in staging ML service"
        else
            warning "GPU not available in staging ML service"
        fi
    else
        error "ML Training Service test failed"
    fi

    success "Staging tests completed"
}

show_staging_status() {
    log "Staging environment status:"

    # Show Docker containers status
    echo
    echo "Docker Containers:"
    docker-compose -f "$COMPOSE_FILE" ps

    # Show resource usage
    echo
    echo "Resource Usage:"
    docker stats --no-stream --format "table {{.Container}}\t{{.CPUPerc}}\t{{.MemUsage}}\t{{.MemPerc}}" $(docker-compose -f "$COMPOSE_FILE" ps -q) 2>/dev/null || true

    # Show network status
    echo
    echo "Network Status:"
    if docker network ls | grep -q "$STAGING_NETWORK"; then
        success "Staging network exists: $STAGING_NETWORK"
    else
        warning "Staging network not found: $STAGING_NETWORK"
    fi

    # Show volume usage
    echo
    echo "Volume Usage:"
    docker volume ls | grep staging || true
}

show_staging_logs() {
    local service="$1"

    if [ -z "$service" ]; then
        log "Available services:"
        docker-compose -f "$COMPOSE_FILE" config --services
        return 1
    fi

    log "Showing logs for staging service: $service"
    docker-compose -f "$COMPOSE_FILE" logs --tail=100 -f "$service"
}

cleanup_staging() {
    log "Cleaning up staging environment..."

    # Stop and remove containers
    docker-compose -f "$COMPOSE_FILE" down --remove-orphans

    # Remove staging volumes (optional)
    read -p "Remove staging data volumes? (y/N): " -n 1 -r
    echo
    if [[ $REPLY =~ ^[Yy]$ ]]; then
        docker volume rm $(docker volume ls -q | grep staging) 2>/dev/null || true
        success "Staging volumes removed"
    fi

    # Clean up staging files
    if [ -d "$STAGING_HOME" ]; then
        read -p "Remove staging directory $STAGING_HOME? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            rm -rf "$STAGING_HOME"
            success "Staging directory removed"
        fi
    fi

    success "Staging cleanup completed"
}

reset_staging() {
    log "Resetting staging environment..."

    cleanup_staging
    setup_staging_environment

    success "Staging environment reset completed"
}

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

main() {
    if [ $# -eq 0 ]; then
        usage
    fi

    local command="$1"
    shift

    # Print banner
    echo "============================================================"
    echo "  Foxhunt HFT Staging Deployment Manager"
    echo "============================================================"
    echo

    case $command in
        deploy)
            if [ $# -eq 0 ]; then
                error "Version is required for deploy command"
                usage
            fi
            deploy_to_staging "$@"
            ;;
        test)
            run_staging_tests
            ;;
        status)
            show_staging_status
            ;;
        logs)
            show_staging_logs "$@"
            ;;
        cleanup)
            cleanup_staging
            ;;
        reset)
            reset_staging
            ;;
        -h|--help)
            usage
            ;;
        *)
            error "Unknown command: $command"
            usage
            ;;
    esac

    log "Staging deployment operation completed"
    log "Log file: $STAGING_LOG"
}

# Execute main function
main "$@"