#!/bin/bash

# ==============================================================================
# Foxhunt HFT System - Hyperparameter Tuning Deployment Script
# ==============================================================================
#
# Purpose: Deploy the hyperparameter tuning system with full validation
#
# Prerequisites:
#   - CUDA-capable GPU (RTX 3050 Ti or better)
#   - Docker with NVIDIA runtime
#   - Training data in test_data/real/databento/ml_training/
#   - PostgreSQL, Redis, MinIO infrastructure
#
# Usage:
#   ./scripts/deploy_tuning.sh [--skip-build] [--skip-smoke-test]
#
# Options:
#   --skip-build       Skip cargo build step (use existing binaries)
#   --skip-smoke-test  Skip smoke test validation
#   --rollback         Rollback to previous version
#
# ==============================================================================

set -e  # Exit on error
set -o pipefail  # Catch errors in pipelines

# ==============================================================================
# Configuration
# ==============================================================================

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
BACKUP_DIR="$PROJECT_ROOT/.deployment_backup"
LOG_FILE="/tmp/foxhunt_tuning_deployment.log"

# Timeouts (seconds)
INFRA_HEALTH_TIMEOUT=60
SERVICE_STARTUP_TIMEOUT=120
SMOKE_TEST_TIMEOUT=300

# Flags
SKIP_BUILD=false
SKIP_SMOKE_TEST=false
ROLLBACK_MODE=false

# Color codes
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
CYAN='\033[0;36m'
NC='\033[0m' # No Color

# ==============================================================================
# Helper Functions
# ==============================================================================

log() {
    echo -e "$(date '+%Y-%m-%d %H:%M:%S') $*" | tee -a "$LOG_FILE"
}

log_info() {
    log "${BLUE}[INFO]${NC} $*"
}

log_success() {
    log "${GREEN}[SUCCESS]${NC} $*"
}

log_warning() {
    log "${YELLOW}[WARNING]${NC} $*"
}

log_error() {
    log "${RED}[ERROR]${NC} $*"
}

section_header() {
    echo ""
    echo "================================================================================"
    log_info "$1"
    echo "================================================================================"
}

check_command() {
    local cmd=$1
    local install_hint=$2

    if ! command -v "$cmd" &> /dev/null; then
        log_error "Required command '$cmd' not found"
        if [ -n "$install_hint" ]; then
            log_info "Install with: $install_hint"
        fi
        return 1
    fi
    return 0
}

wait_for_service() {
    local service_name=$1
    local health_check=$2
    local timeout=$3
    local elapsed=0

    log_info "Waiting for $service_name to be healthy (timeout: ${timeout}s)..."

    while [ $elapsed -lt $timeout ]; do
        if eval "$health_check" &> /dev/null; then
            log_success "$service_name is healthy"
            return 0
        fi
        sleep 2
        elapsed=$((elapsed + 2))
        echo -n "."
    done

    echo ""
    log_error "$service_name failed to become healthy within ${timeout}s"
    return 1
}

# ==============================================================================
# Backup and Rollback
# ==============================================================================

create_backup() {
    section_header "Creating Deployment Backup"

    mkdir -p "$BACKUP_DIR"

    # Backup Docker images
    log_info "Backing up Docker images..."
    docker images --format "{{.Repository}}:{{.Tag}}" | grep "foxhunt" > "$BACKUP_DIR/images.txt" || true

    # Backup .env file
    if [ -f "$PROJECT_ROOT/.env" ]; then
        cp "$PROJECT_ROOT/.env" "$BACKUP_DIR/.env.backup"
        log_success "Backed up .env file"
    fi

    # Save current git commit
    cd "$PROJECT_ROOT"
    git rev-parse HEAD > "$BACKUP_DIR/git_commit.txt" 2>/dev/null || echo "unknown" > "$BACKUP_DIR/git_commit.txt"

    log_success "Backup created in $BACKUP_DIR"
}

rollback_deployment() {
    section_header "Rolling Back Deployment"

    if [ ! -d "$BACKUP_DIR" ]; then
        log_error "No backup found at $BACKUP_DIR"
        exit 1
    fi

    # Stop current services
    log_info "Stopping current services..."
    docker-compose down || true

    # Restore .env
    if [ -f "$BACKUP_DIR/.env.backup" ]; then
        cp "$BACKUP_DIR/.env.backup" "$PROJECT_ROOT/.env"
        log_success "Restored .env file"
    fi

    # Restore git commit
    if [ -f "$BACKUP_DIR/git_commit.txt" ]; then
        local commit=$(cat "$BACKUP_DIR/git_commit.txt")
        if [ "$commit" != "unknown" ]; then
            log_info "Previous commit was: $commit"
            log_warning "Run 'git checkout $commit' to restore code"
        fi
    fi

    log_success "Rollback preparation complete"
    log_info "Restart services with: docker-compose up -d"
}

# ==============================================================================
# Prerequisites Check
# ==============================================================================

check_prerequisites() {
    section_header "Checking Prerequisites"

    local failed=false

    # 1. Check required commands
    log_info "Checking required commands..."
    check_command "docker" "See: https://docs.docker.com/get-docker/" || failed=true
    check_command "docker-compose" "See: https://docs.docker.com/compose/install/" || failed=true
    check_command "cargo" "See: https://rustup.rs/" || failed=true
    check_command "psql" "apt-get install postgresql-client" || failed=true
    check_command "redis-cli" "apt-get install redis-tools" || failed=true

    # 2. Check CUDA availability
    log_info "Checking CUDA availability..."
    if command -v nvidia-smi &> /dev/null; then
        log_success "CUDA available:"
        nvidia-smi --query-gpu=name,driver_version,memory.total --format=csv,noheader | head -1
    else
        log_error "nvidia-smi not found - GPU acceleration unavailable"
        failed=true
    fi

    # 3. Check Docker daemon
    log_info "Checking Docker daemon..."
    if ! docker info &> /dev/null; then
        log_error "Docker daemon not running"
        failed=true
    else
        log_success "Docker daemon running"
    fi

    # 4. Check Docker Compose version
    log_info "Checking Docker Compose version..."
    local compose_version=$(docker-compose version --short 2>/dev/null || echo "0.0.0")
    local major_version=$(echo "$compose_version" | cut -d. -f1)
    if [ "$major_version" -ge 2 ]; then
        log_success "Docker Compose version: $compose_version (✓ ≥ 2.0)"
    else
        log_error "Docker Compose version $compose_version < 2.0"
        failed=true
    fi

    # 5. Check NVIDIA Docker runtime
    log_info "Checking NVIDIA Docker runtime..."
    if docker run --rm --gpus all nvidia/cuda:12.0.0-base-ubuntu22.04 nvidia-smi &> /dev/null; then
        log_success "NVIDIA Docker runtime available"
    else
        log_error "NVIDIA Docker runtime not available"
        log_info "Install with: distribution=$(. /etc/os-release;echo \$ID\$VERSION_ID) && curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add - && curl -s -L https://nvidia.github.io/nvidia-docker/\$distribution/nvidia-docker.list | sudo tee /etc/apt/sources.list.d/nvidia-docker.list && sudo apt-get update && sudo apt-get install -y nvidia-docker2 && sudo systemctl restart docker"
        failed=true
    fi

    # 6. Check training data
    log_info "Checking training data..."
    local data_dir="$PROJECT_ROOT/test_data/real/databento/ml_training"
    if [ -d "$data_dir" ]; then
        local file_count=$(find "$data_dir" -name "*.dbn" | wc -l)
        if [ "$file_count" -gt 0 ]; then
            log_success "Found $file_count DBN training files"
        else
            log_error "No *.dbn files found in $data_dir"
            failed=true
        fi
    else
        log_error "Training data directory not found: $data_dir"
        failed=true
    fi

    # 7. Check .env file
    log_info "Checking .env configuration..."
    if [ ! -f "$PROJECT_ROOT/.env" ]; then
        log_warning ".env file not found, using defaults from docker-compose.yml"
        log_info "For production, copy .env.example to .env and configure"
    else
        log_success ".env file found"

        # Check JWT_SECRET
        if grep -q "JWT_SECRET=your_jwt_secret_here_change_in_production" "$PROJECT_ROOT/.env"; then
            log_warning "JWT_SECRET is using default value - change for production!"
        fi
    fi

    if [ "$failed" = true ]; then
        log_error "Prerequisites check failed"
        exit 1
    fi

    log_success "All prerequisites met"
}

# ==============================================================================
# Infrastructure Deployment
# ==============================================================================

deploy_infrastructure() {
    section_header "Deploying Infrastructure Services"

    cd "$PROJECT_ROOT"

    # Start infrastructure services
    log_info "Starting PostgreSQL, Redis, MinIO..."
    docker-compose up -d postgres redis minio

    # Wait for PostgreSQL
    wait_for_service "PostgreSQL" \
        "docker exec foxhunt-postgres pg_isready -U foxhunt" \
        "$INFRA_HEALTH_TIMEOUT" || exit 1

    # Wait for Redis
    wait_for_service "Redis" \
        "docker exec foxhunt-redis redis-cli ping | grep -q PONG" \
        "$INFRA_HEALTH_TIMEOUT" || exit 1

    # Wait for MinIO
    wait_for_service "MinIO" \
        "curl -sf http://localhost:9000/minio/health/live" \
        "$INFRA_HEALTH_TIMEOUT" || exit 1

    # Initialize MinIO bucket
    log_info "Initializing MinIO bucket..."
    docker run --rm --network foxhunt-network \
        -e MC_HOST_minio=http://foxhunt:foxhunt_dev_password@minio:9000 \
        minio/mc:latest \
        mb minio/ml-models --ignore-existing || true

    log_success "Infrastructure services healthy"
}

# ==============================================================================
# Build Step
# ==============================================================================

build_services() {
    section_header "Building Services"

    cd "$PROJECT_ROOT"

    if [ "$SKIP_BUILD" = true ]; then
        log_warning "Skipping build step (--skip-build flag)"
        return 0
    fi

    # Build ml crate with CUDA
    log_info "Building ml crate with CUDA support..."
    cargo build -p ml --release --features cuda 2>&1 | tee -a "$LOG_FILE"

    if [ ${PIPESTATUS[0]} -ne 0 ]; then
        log_error "ML crate build failed"
        exit 1
    fi
    log_success "ML crate built successfully"

    # Build ML Training Service Docker image
    log_info "Building ML Training Service Docker image..."
    docker-compose build ml_training_service 2>&1 | tee -a "$LOG_FILE"

    if [ ${PIPESTATUS[0]} -ne 0 ]; then
        log_error "ML Training Service Docker build failed"
        exit 1
    fi
    log_success "ML Training Service Docker image built"

    # Build API Gateway Docker image
    log_info "Building API Gateway Docker image..."
    docker-compose build api_gateway 2>&1 | tee -a "$LOG_FILE"

    if [ ${PIPESTATUS[0]} -ne 0 ]; then
        log_error "API Gateway Docker build failed"
        exit 1
    fi
    log_success "API Gateway Docker image built"

    log_success "All services built successfully"
}

# ==============================================================================
# Service Deployment
# ==============================================================================

deploy_services() {
    section_header "Deploying Application Services"

    cd "$PROJECT_ROOT"

    # Start ML Training Service
    log_info "Starting ML Training Service..."
    docker-compose up -d ml_training_service

    # Wait for ML service (longer timeout for model loading)
    wait_for_service "ML Training Service" \
        "curl -sf http://localhost:8095/health" \
        "$SERVICE_STARTUP_TIMEOUT" || exit 1

    # Start API Gateway
    log_info "Starting API Gateway..."
    docker-compose up -d api_gateway

    # Wait for API Gateway
    wait_for_service "API Gateway" \
        "docker exec foxhunt-api-gateway /usr/local/bin/grpc_health_probe -addr=localhost:50050 2>/dev/null || curl -sf http://localhost:8080/health" \
        60 || exit 1

    log_success "All services deployed and healthy"
}

# ==============================================================================
# Smoke Tests
# ==============================================================================

run_smoke_tests() {
    section_header "Running Smoke Tests"

    if [ "$SKIP_SMOKE_TEST" = true ]; then
        log_warning "Skipping smoke tests (--skip-smoke-test flag)"
        return 0
    fi

    cd "$PROJECT_ROOT"

    # 1. Check PostgreSQL connectivity
    log_info "Testing PostgreSQL connectivity..."
    if psql postgresql://foxhunt:foxhunt_dev_password@localhost:5432/foxhunt -c '\dt' &> /dev/null; then
        log_success "PostgreSQL connectivity OK"
    else
        log_error "PostgreSQL connectivity failed"
        exit 1
    fi

    # 2. Check Redis connectivity
    log_info "Testing Redis connectivity..."
    if redis-cli -h localhost -p 6379 PING | grep -q PONG; then
        log_success "Redis connectivity OK"
    else
        log_error "Redis connectivity failed"
        exit 1
    fi

    # 3. Check MinIO connectivity
    log_info "Testing MinIO connectivity..."
    if curl -sf http://localhost:9000/minio/health/live &> /dev/null; then
        log_success "MinIO connectivity OK"
    else
        log_error "MinIO connectivity failed"
        exit 1
    fi

    # 4. Test TLI auth (if TLI is built)
    log_info "Testing TLI authentication..."
    if [ -f "$PROJECT_ROOT/target/release/tli" ] || [ -f "$PROJECT_ROOT/target/debug/tli" ]; then
        # Note: This requires TLI to be configured with test credentials
        # For now, just check if the binary exists
        log_success "TLI binary found (auth test requires manual configuration)"
    else
        log_warning "TLI binary not found - skipping auth test"
        log_info "Build TLI with: cargo build -p tli --release"
    fi

    # 5. Test hyperparameter tuning API (single trial)
    log_info "Testing hyperparameter tuning with single trial..."

    # Create a test tuning request
    local test_request=$(cat <<EOF
{
  "model_type": "DQN",
  "n_trials": 1,
  "timeout": 60,
  "search_space": {
    "learning_rate": {
      "type": "loguniform",
      "low": 1e-5,
      "high": 1e-3
    },
    "batch_size": {
      "type": "categorical",
      "choices": [32, 64]
    }
  }
}
EOF
    )

    # Note: This requires gRPC client tools or TLI
    # For now, just verify the service is listening
    if docker exec foxhunt-ml-training-service /usr/local/bin/grpc_health_probe -addr=localhost:50053 &> /dev/null; then
        log_success "ML Training Service gRPC endpoint responding"
    else
        log_error "ML Training Service gRPC endpoint not responding"
        exit 1
    fi

    # 6. Check MinIO for study storage
    log_info "Verifying MinIO bucket setup..."
    if docker run --rm --network foxhunt-network \
        -e MC_HOST_minio=http://foxhunt:foxhunt_dev_password@minio:9000 \
        minio/mc:latest \
        ls minio/ml-models &> /dev/null; then
        log_success "MinIO ml-models bucket accessible"
    else
        log_error "MinIO ml-models bucket not accessible"
        exit 1
    fi

    # 7. Test service health endpoints
    log_info "Testing service health endpoints..."

    local services=(
        "API Gateway:http://localhost:8080/health"
        "ML Training Service:http://localhost:8095/health"
    )

    for service_info in "${services[@]}"; do
        local name="${service_info%%:*}"
        local url="${service_info#*:}"

        if curl -sf "$url" &> /dev/null; then
            log_success "$name health check OK"
        else
            log_error "$name health check failed"
            exit 1
        fi
    done

    log_success "All smoke tests passed"
}

# ==============================================================================
# Post-Deployment Validation
# ==============================================================================

validate_deployment() {
    section_header "Post-Deployment Validation"

    cd "$PROJECT_ROOT"

    # Show service status
    log_info "Service status:"
    docker-compose ps

    echo ""
    log_info "Container logs (last 10 lines):"
    for service in ml_training_service api_gateway; do
        echo ""
        echo "--- $service ---"
        docker-compose logs --tail=10 "$service" 2>&1 | tail -10
    done

    # Check GPU usage
    echo ""
    log_info "GPU status:"
    nvidia-smi --query-gpu=name,utilization.gpu,memory.used,memory.total --format=csv,noheader

    # Summary
    echo ""
    section_header "Deployment Summary"

    log_success "Hyperparameter tuning system deployed successfully!"
    echo ""
    echo "Service Endpoints:"
    echo "  - API Gateway (gRPC):      localhost:50051"
    echo "  - ML Training Service:     localhost:50054"
    echo "  - PostgreSQL:              localhost:5432"
    echo "  - Redis:                   localhost:6379"
    echo "  - MinIO:                   localhost:9000 (console: 9001)"
    echo ""
    echo "Web UIs:"
    echo "  - MinIO Console:           http://localhost:9001 (foxhunt/foxhunt_dev_password)"
    echo ""
    echo "Next Steps:"
    echo "  1. Test tuning via TLI:    tli tune start --model DQN --trials 10"
    echo "  2. Check tuning status:    tli tune status"
    echo "  3. View best parameters:   tli tune best"
    echo "  4. Monitor GPU usage:      watch -n 1 nvidia-smi"
    echo "  5. View service logs:      docker-compose logs -f ml_training_service"
    echo ""
    echo "Configuration:"
    echo "  - Study storage:           PostgreSQL (Optuna)"
    echo "  - Model checkpoints:       MinIO (S3-compatible)"
    echo "  - Training data:           test_data/real/databento/ml_training/"
    echo ""
    echo "Troubleshooting:"
    echo "  - View logs:               docker-compose logs ml_training_service"
    echo "  - Restart service:         docker-compose restart ml_training_service"
    echo "  - Rollback deployment:     ./scripts/deploy_tuning.sh --rollback"
    echo ""
}

# ==============================================================================
# Main Execution
# ==============================================================================

main() {
    # Parse arguments
    while [ $# -gt 0 ]; do
        case "$1" in
            --skip-build)
                SKIP_BUILD=true
                shift
                ;;
            --skip-smoke-test)
                SKIP_SMOKE_TEST=true
                shift
                ;;
            --rollback)
                ROLLBACK_MODE=true
                shift
                ;;
            --help|-h)
                cat << EOF
Usage: $0 [OPTIONS]

Deploy the Foxhunt hyperparameter tuning system.

OPTIONS:
    --skip-build        Skip cargo build step (use existing binaries)
    --skip-smoke-test   Skip smoke test validation
    --rollback          Rollback to previous version
    --help, -h          Show this help message

EXAMPLES:
    # Full deployment
    $0

    # Quick deployment (skip rebuild)
    $0 --skip-build

    # Deployment without smoke tests
    $0 --skip-smoke-test

    # Rollback to previous version
    $0 --rollback

EOF
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                log_info "Use --help for usage information"
                exit 1
                ;;
        esac
    done

    # Start deployment
    echo "================================================================================"
    echo "         Foxhunt HFT System - Hyperparameter Tuning Deployment"
    echo "================================================================================"
    echo "Started at: $(date '+%Y-%m-%d %H:%M:%S')"
    echo "Log file: $LOG_FILE"
    echo ""

    # Initialize log file
    echo "=== Foxhunt Tuning Deployment Log ===" > "$LOG_FILE"
    echo "Started: $(date)" >> "$LOG_FILE"
    echo "" >> "$LOG_FILE"

    # Handle rollback mode
    if [ "$ROLLBACK_MODE" = true ]; then
        rollback_deployment
        exit 0
    fi

    # Deployment steps
    trap 'log_error "Deployment failed! Check logs at $LOG_FILE"; exit 1' ERR

    check_prerequisites
    create_backup
    deploy_infrastructure
    build_services
    deploy_services
    run_smoke_tests
    validate_deployment

    # Success
    log_success "Deployment completed successfully at $(date '+%Y-%m-%d %H:%M:%S')"
    log_info "Full deployment log available at: $LOG_FILE"
}

# Run main function
main "$@"