#!/bin/bash
# Canary traffic splitting configuration for Foxhunt HFT Trading System
# Implements precise traffic percentage routing with health monitoring

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
DEPLOYMENT_LOG="/home/jgrusewski/Work/foxhunt/logs/canary-traffic-$(date +%s).log"
NGINX_CONF_DIR="/etc/nginx/conf.d"
FOXHUNT_UPSTREAM_CONF="$NGINX_CONF_DIR/foxhunt-upstream.conf"
CANARY_CONF="$NGINX_CONF_DIR/foxhunt-canary.conf"
MAIN_CONF="$NGINX_CONF_DIR/foxhunt-main.conf"

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

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

error() {
    echo -e "${RED}ERROR: $1${NC}" | tee -a "$DEPLOYMENT_LOG"
}

success() {
    echo -e "${GREEN}SUCCESS: $1${NC}" | tee -a "$DEPLOYMENT_LOG"
}

warning() {
    echo -e "${YELLOW}WARNING: $1${NC}" | tee -a "$DEPLOYMENT_LOG"
}

info() {
    echo -e "${BLUE}INFO: $1${NC}" | tee -a "$DEPLOYMENT_LOG"
}

usage() {
    echo "Usage: $0 <canary_percentage> [options]"
    echo ""
    echo "Arguments:"
    echo "  canary_percentage    Percentage of traffic to route to canary (1-99)"
    echo ""
    echo "Options:"
    echo "  --canary-port PORT   Base port for canary services (default: 8085)"
    echo "  --main-port PORT     Base port for main services (default: 8080)"
    echo "  --validate-only      Only validate configuration, don't apply"
    echo "  --remove-canary      Remove canary configuration and route all traffic to main"
    echo "  --health-check       Perform health check before applying configuration"
    echo "  -h, --help          Show this help message"
    echo ""
    echo "Examples:"
    echo "  $0 1                 # Route 1% traffic to canary, 99% to main"
    echo "  $0 5 --canary-port 8085"
    echo "  $0 --remove-canary   # Remove canary routing"
    exit 1
}

validate_percentage() {
    local percentage="$1"
    
    if ! [[ "$percentage" =~ ^[0-9]+$ ]] || [ "$percentage" -lt 1 ] || [ "$percentage" -gt 99 ]; then
        error "Invalid percentage: $percentage. Must be between 1 and 99"
        return 1
    fi
    
    return 0
}

health_check_services() {
    local main_port="$1"
    local canary_port="$2"
    
    log "Performing health checks..."
    
    # Check main services
    local services=("core" "tli" "ml" "risk" "data")
    
    for i in "${!services[@]}"; do
        local service="${services[$i]}"
        local main_health_url="http://localhost:$((main_port + i))/health"
        local canary_health_url="http://localhost:$((canary_port + i))/health"
        
        # Check main service
        if ! curl -f -s "$main_health_url" > /dev/null 2>&1; then
            error "Main $service service health check failed at $main_health_url"
            return 1
        fi
        
        # Check canary service
        if ! curl -f -s "$canary_health_url" > /dev/null 2>&1; then
            error "Canary $service service health check failed at $canary_health_url"
            return 1
        fi
    done
    
    success "All services passed health checks"
    return 0
}

generate_upstream_config() {
    local canary_percentage="$1"
    local main_port="$2"
    local canary_port="$3"
    
    log "Generating upstream configuration for ${canary_percentage}% canary traffic..."
    
    # Calculate weights for nginx upstream
    # nginx uses weight-based load balancing
    local main_weight=$((100 - canary_percentage))
    local canary_weight="$canary_percentage"
    
    # For very small percentages, we need to scale up to ensure precision
    if [ "$canary_percentage" -lt 5 ]; then
        local scale_factor=20
        main_weight=$((main_weight * scale_factor))
        canary_weight=$((canary_weight * scale_factor))
    fi
    
    cat > "$FOXHUNT_UPSTREAM_CONF" << EOF
# Foxhunt HFT Trading System - Canary Traffic Configuration
# Generated on $(date)
# Main traffic: ${main_weight}, Canary traffic: ${canary_weight}

upstream foxhunt_core {
    server 127.0.0.1:${main_port} weight=${main_weight} max_fails=2 fail_timeout=30s;
    server 127.0.0.1:${canary_port} weight=${canary_weight} max_fails=2 fail_timeout=30s;
}

upstream foxhunt_tli {
    server 127.0.0.1:$((main_port + 1)) weight=${main_weight} max_fails=2 fail_timeout=30s;
    server 127.0.0.1:$((canary_port + 1)) weight=${canary_weight} max_fails=2 fail_timeout=30s;
}

upstream foxhunt_ml {
    server 127.0.0.1:$((main_port + 2)) weight=${main_weight} max_fails=2 fail_timeout=30s;
    server 127.0.0.1:$((canary_port + 2)) weight=${canary_weight} max_fails=2 fail_timeout=30s;
}

upstream foxhunt_risk {
    server 127.0.0.1:$((main_port + 3)) weight=${main_weight} max_fails=2 fail_timeout=30s;
    server 127.0.0.1:$((canary_port + 3)) weight=${canary_weight} max_fails=2 fail_timeout=30s;
}

upstream foxhunt_data {
    server 127.0.0.1:$((main_port + 4)) weight=${main_weight} max_fails=2 fail_timeout=30s;
    server 127.0.0.1:$((canary_port + 4)) weight=${canary_weight} max_fails=2 fail_timeout=30s;
}

# Real-time monitoring endpoint for canary traffic
upstream foxhunt_canary_monitor {
    server 127.0.0.1:9099;  # Dedicated monitoring service
}
EOF
    
    success "Upstream configuration generated with ${canary_percentage}% canary traffic"
}

generate_main_config() {
    local main_port="$1"
    
    log "Generating main-only upstream configuration..."
    
    cat > "$FOXHUNT_UPSTREAM_CONF" << EOF
# Foxhunt HFT Trading System - Main Traffic Configuration
# Generated on $(date)
# All traffic routed to main services

upstream foxhunt_core {
    server 127.0.0.1:${main_port} max_fails=2 fail_timeout=30s;
}

upstream foxhunt_tli {
    server 127.0.0.1:$((main_port + 1)) max_fails=2 fail_timeout=30s;
}

upstream foxhunt_ml {
    server 127.0.0.1:$((main_port + 2)) max_fails=2 fail_timeout=30s;
}

upstream foxhunt_risk {
    server 127.0.0.1:$((main_port + 3)) max_fails=2 fail_timeout=30s;
}

upstream foxhunt_data {
    server 127.0.0.1:$((main_port + 4)) max_fails=2 fail_timeout=30s;
}
EOF
    
    success "Main-only upstream configuration generated"
}

generate_canary_monitoring() {
    local canary_percentage="$1"
    
    log "Generating canary monitoring configuration..."
    
    cat > "$CANARY_CONF" << EOF
# Foxhunt Canary Monitoring Configuration
server {
    listen 9099;
    server_name localhost;
    
    location /canary/status {
        return 200 '{"canary_percentage": ${canary_percentage}, "status": "active", "timestamp": "${$(date -Iseconds)}"}';
        add_header Content-Type application/json;
    }
    
    location /canary/metrics {
        stub_status on;
        access_log off;
    }
    
    # Real-time canary metrics
    location /canary/traffic {
        content_by_lua_block {
            local canary_pct = ${canary_percentage}
            local main_pct = 100 - canary_pct
            
            ngx.header.content_type = "application/json"
            ngx.say('{"main_traffic_pct": ' .. main_pct .. ', "canary_traffic_pct": ' .. canary_pct .. ', "active": true}')
        }
    }
}
EOF
    
    success "Canary monitoring configuration generated"
}

apply_configuration() {
    log "Applying nginx configuration..."
    
    # Test nginx configuration
    if ! sudo nginx -t; then
        error "Nginx configuration test failed"
        return 1
    fi
    
    # Reload nginx gracefully
    if ! sudo systemctl reload nginx; then
        error "Nginx reload failed"
        return 1
    fi
    
    success "Nginx configuration applied successfully"
    
    # Verify configuration is active
    sleep 2
    
    if curl -f -s http://localhost:9099/canary/status > /dev/null 2>&1; then
        info "Canary monitoring endpoint is active"
    else
        warning "Canary monitoring endpoint not responding"
    fi
    
    return 0
}

remove_canary_config() {
    local main_port="$1"
    
    log "Removing canary configuration..."
    
    # Generate main-only config
    generate_main_config "$main_port"
    
    # Remove canary-specific configs
    sudo rm -f "$CANARY_CONF"
    
    # Apply configuration
    apply_configuration
    
    success "Canary configuration removed, all traffic routed to main services"
}

monitor_traffic_split() {
    local duration=60  # Monitor for 60 seconds
    local interval=10  # Check every 10 seconds
    
    log "Monitoring traffic distribution for ${duration} seconds..."
    
    for i in $(seq 0 $interval $((duration - interval))); do
        # Get access log entries from the last interval
        local main_requests
        local canary_requests
        
        main_requests=$(sudo tail -n 1000 /var/log/nginx/access.log | grep -c "127.0.0.1:8080" || echo "0")
        canary_requests=$(sudo tail -n 1000 /var/log/nginx/access.log | grep -c "127.0.0.1:8085" || echo "0")
        
        local total_requests=$((main_requests + canary_requests))
        
        if [ "$total_requests" -gt 0 ]; then
            local actual_canary_pct=$((canary_requests * 100 / total_requests))
            info "Traffic distribution - Main: $((100 - actual_canary_pct))%, Canary: ${actual_canary_pct}% (${total_requests} total requests)"
        else
            info "No traffic observed in last ${interval} seconds"
        fi
        
        sleep "$interval"
    done
}

main() {
    local canary_percentage=""
    local canary_port=8085
    local main_port=8080
    local validate_only=false
    local remove_canary=false
    local health_check=false
    
    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --canary-port)
                canary_port="$2"
                shift 2
                ;;
            --main-port)
                main_port="$2"
                shift 2
                ;;
            --validate-only)
                validate_only=true
                shift
                ;;
            --remove-canary)
                remove_canary=true
                shift
                ;;
            --health-check)
                health_check=true
                shift
                ;;
            -h|--help)
                usage
                ;;
            *)
                if [ -z "$canary_percentage" ] && [[ "$1" =~ ^[0-9]+$ ]]; then
                    canary_percentage="$1"
                else
                    error "Unknown option: $1"
                    usage
                fi
                shift
                ;;
        esac
    done
    
    log "Starting canary traffic configuration..."
    log "Main port base: $main_port"
    log "Canary port base: $canary_port"
    log "Validate only: $validate_only"
    log "Remove canary: $remove_canary"
    log "Health check: $health_check"
    
    # Handle canary removal
    if [ "$remove_canary" == "true" ]; then
        if [ "$validate_only" == "true" ]; then
            log "Validation-only mode: would remove canary configuration"
            return 0
        fi
        
        remove_canary_config "$main_port"
        return $?
    fi
    
    # Validate percentage
    if [ -z "$canary_percentage" ]; then
        error "Canary percentage is required"
        usage
    fi
    
    if ! validate_percentage "$canary_percentage"; then
        return 1
    fi
    
    log "Configuring ${canary_percentage}% canary traffic routing..."
    
    # Health check if requested
    if [ "$health_check" == "true" ]; then
        if ! health_check_services "$main_port" "$canary_port"; then
            error "Health check failed, aborting configuration"
            return 1
        fi
    fi
    
    # Validation only mode
    if [ "$validate_only" == "true" ]; then
        log "Validation-only mode: would configure ${canary_percentage}% canary traffic"
        return 0
    fi
    
    # Generate configurations
    generate_upstream_config "$canary_percentage" "$main_port" "$canary_port"
    generate_canary_monitoring "$canary_percentage"
    
    # Apply configuration
    if ! apply_configuration; then
        error "Failed to apply configuration"
        return 1
    fi
    
    success "Canary traffic configuration completed"
    info "Traffic distribution: Main ${$((100 - canary_percentage))}%, Canary ${canary_percentage}%"
    
    # Optional traffic monitoring
    monitor_traffic_split
    
    return 0
}

# Execute main function
main "$@"