#!/bin/bash

# Foxhunt HFT System - TLI Tuning Workflow Test Script
# Tests the full hyperparameter tuning workflow from start to completion

set -euo pipefail

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

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "$SCRIPT_DIR/.." && pwd)"
TLI_BIN="$PROJECT_ROOT/target/release/tli"
API_GATEWAY_URL="http://localhost:50051"
ML_SERVICE_URL="http://localhost:8095"
POLL_INTERVAL=5
MAX_WAIT_TIME=300  # 5 minutes max wait
JWT_TOKEN_FILE="$HOME/.foxhunt/jwt_token"

# Test parameters
MODEL_TYPE="DQN"
NUM_TRIALS=5
CONFIG_FILE="$PROJECT_ROOT/test_data/tuning_config.yaml"
DATA_SOURCE="$PROJECT_ROOT/test_data/btcusdt_sample_100.parquet"

# State tracking
JOB_ID=""
START_TIME=$(date +%s)

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

print_banner() {
    echo ""
    echo "================================================================================"
    echo "  $1"
    echo "================================================================================"
    echo ""
}

print_step() {
    echo -e "${BLUE}[STEP]${NC} $1"
}

print_success() {
    echo -e "${GREEN}✓${NC} $1"
}

print_error() {
    echo -e "${RED}✗${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}⚠${NC} $1"
}

print_info() {
    echo -e "${CYAN}ℹ${NC} $1"
}

elapsed_time() {
    local now=$(date +%s)
    echo $((now - START_TIME))
}

format_duration() {
    local seconds=$1
    if [ $seconds -lt 60 ]; then
        echo "${seconds}s"
    else
        local minutes=$((seconds / 60))
        local secs=$((seconds % 60))
        echo "${minutes}m ${secs}s"
    fi
}

cleanup() {
    local exit_code=$?

    echo ""
    if [ $exit_code -ne 0 ]; then
        print_error "Test failed with exit code $exit_code"

        # Stop job if it was started
        if [ -n "$JOB_ID" ]; then
            print_info "Attempting to stop job $JOB_ID..."
            stop_tuning_job "Test script cleanup"
        fi
    fi

    local total_time=$(elapsed_time)
    print_info "Total test duration: $(format_duration $total_time)"
    exit $exit_code
}

trap cleanup EXIT INT TERM

# ============================================================================
# Prerequisite Checks
# ============================================================================

check_prerequisites() {
    print_banner "Prerequisite Checks"

    local all_ok=true

    # Check TLI binary exists
    print_step "Checking TLI binary..."
    if [ ! -f "$TLI_BIN" ]; then
        print_error "TLI binary not found at: $TLI_BIN"
        print_info "Build with: cargo build --release -p tli"
        all_ok=false
    else
        print_success "TLI binary found"
        print_info "Location: $TLI_BIN"
    fi

    # Check JWT token exists
    print_step "Checking JWT token..."
    if [ ! -f "$JWT_TOKEN_FILE" ]; then
        print_error "JWT token not found at: $JWT_TOKEN_FILE"
        print_info "Authenticate with: $TLI_BIN auth login"
        all_ok=false
    else
        local jwt_token=$(cat "$JWT_TOKEN_FILE")
        if [ -z "$jwt_token" ]; then
            print_error "JWT token file is empty"
            all_ok=false
        else
            print_success "JWT token found"
            # Mask token for security
            local masked_token="${jwt_token:0:20}...${jwt_token: -10}"
            print_info "Token (masked): $masked_token"
        fi
    fi

    # Check API Gateway is running
    print_step "Checking API Gateway health..."
    if ! curl -sf http://localhost:8080/health > /dev/null 2>&1; then
        print_error "API Gateway not responding at http://localhost:8080/health"
        print_info "Start services with: docker-compose up -d"
        all_ok=false
    else
        print_success "API Gateway is healthy"
    fi

    # Check ML Training Service is running
    print_step "Checking ML Training Service health..."
    if ! curl -sf "$ML_SERVICE_URL/health" > /dev/null 2>&1; then
        print_error "ML Training Service not responding at $ML_SERVICE_URL/health"
        print_info "Start with: cargo run --release -p ml_training_service"
        all_ok=false
    else
        print_success "ML Training Service is healthy"
    fi

    # Check config file exists
    print_step "Checking tuning config file..."
    if [ ! -f "$CONFIG_FILE" ]; then
        print_warning "Config file not found, will create minimal config"
        create_minimal_config
    else
        print_success "Config file found"
        print_info "Path: $CONFIG_FILE"
    fi

    # Check test data exists
    print_step "Checking test data..."
    if [ ! -f "$DATA_SOURCE" ]; then
        print_warning "Test data not found at: $DATA_SOURCE"
        print_info "Will use default data source from config"
        DATA_SOURCE=""
    else
        print_success "Test data found"
        print_info "Path: $DATA_SOURCE"

        # Show file size
        local file_size=$(du -h "$DATA_SOURCE" | cut -f1)
        print_info "Size: $file_size"
    fi

    # Check grpcurl if available (optional)
    if command -v grpcurl &> /dev/null; then
        print_success "grpcurl available for direct gRPC testing"
    else
        print_warning "grpcurl not installed (optional)"
        print_info "Install: go install github.com/fullstorydev/grpcurl/cmd/grpcurl@latest"
    fi

    echo ""

    if [ "$all_ok" = false ]; then
        print_error "Prerequisites check FAILED"
        print_info "Please resolve issues above and retry"
        exit 1
    else
        print_success "All prerequisites satisfied"
    fi
}

create_minimal_config() {
    local config_dir=$(dirname "$CONFIG_FILE")
    mkdir -p "$config_dir"

    cat > "$CONFIG_FILE" << 'EOF'
# Minimal tuning configuration for testing
# Generated by test_tli_tuning.sh

tuning:
  # Optuna study configuration
  study_name: "test_study"
  storage: "sqlite:///optuna_test.db"
  direction: "maximize"  # Maximize Sharpe ratio

  # Search space for DQN
  search_space:
    learning_rate:
      type: "float"
      low: 0.0001
      high: 0.01
      log: true

    batch_size:
      type: "int"
      low: 32
      high: 128
      step: 32

    gamma:
      type: "float"
      low: 0.90
      high: 0.99

    epsilon_decay:
      type: "float"
      low: 0.990
      high: 0.999

  # Training configuration
  training:
    epochs: 10
    validation_split: 0.2
    early_stopping_patience: 3

  # Evaluation metrics
  metrics:
    - "sharpe_ratio"
    - "total_return"
    - "max_drawdown"
    - "win_rate"
EOF

    print_success "Created minimal config: $CONFIG_FILE"
}

# ============================================================================
# Tuning Workflow Functions
# ============================================================================

start_tuning_job() {
    print_banner "Starting Tuning Job"

    print_step "Submitting tuning job to TLI..."
    echo ""

    # Build TLI command
    local cmd="$TLI_BIN tune start --model $MODEL_TYPE --trials $NUM_TRIALS --config $CONFIG_FILE"
    if [ -n "$DATA_SOURCE" ]; then
        cmd="$cmd --data-source $DATA_SOURCE"
    fi

    print_info "Command: $cmd"
    echo ""

    # Execute and capture output
    local output
    if ! output=$($cmd 2>&1); then
        print_error "Failed to start tuning job"
        echo "$output"
        exit 1
    fi

    echo "$output"
    echo ""

    # Extract job ID from output
    JOB_ID=$(echo "$output" | grep -oP 'Job ID: \K[0-9a-f-]+' || true)

    if [ -z "$JOB_ID" ]; then
        print_error "Failed to extract job ID from output"
        exit 1
    fi

    print_success "Tuning job started successfully"
    print_info "Job ID: ${MAGENTA}${JOB_ID}${NC}"

    # Save to tracking file
    local tracking_file="$HOME/.foxhunt/test_tuning_job.txt"
    echo "$JOB_ID" > "$tracking_file"
    print_info "Saved job ID to: $tracking_file"
}

poll_status() {
    print_banner "Polling Job Status"

    local wait_time=0
    local previous_progress=""

    print_info "Polling every ${POLL_INTERVAL}s (max wait: $(format_duration $MAX_WAIT_TIME))"
    echo ""

    while [ $wait_time -lt $MAX_WAIT_TIME ]; do
        # Get status
        local output
        if ! output=$($TLI_BIN tune status --job-id "$JOB_ID" 2>&1); then
            print_error "Failed to get job status"
            echo "$output"
            exit 1
        fi

        # Extract status line
        local status=$(echo "$output" | grep -oP 'Status: \K\w+' || echo "UNKNOWN")
        local progress=$(echo "$output" | grep -oP 'Progress: \K[0-9]+/[0-9]+' || echo "0/0")
        local percent=$(echo "$output" | grep -oP '\(([0-9.]+)%\)' | grep -oP '[0-9.]+' || echo "0.0")

        # Only print if progress changed
        if [ "$progress" != "$previous_progress" ]; then
            local timestamp=$(date '+%H:%M:%S')
            echo -e "${CYAN}[$timestamp]${NC} Status: ${YELLOW}$status${NC} | Progress: ${GREEN}$progress${NC} (${percent}%) | Elapsed: $(format_duration $wait_time)"
            previous_progress="$progress"
        fi

        # Check if completed, failed, or stopped
        case "$status" in
            "TUNING_COMPLETED"|"COMPLETED")
                echo ""
                print_success "Job completed successfully!"
                echo ""
                echo "$output"
                return 0
                ;;
            "TUNING_FAILED"|"FAILED")
                echo ""
                print_error "Job failed!"
                echo ""
                echo "$output"
                exit 1
                ;;
            "TUNING_STOPPED"|"STOPPED")
                echo ""
                print_warning "Job was stopped"
                echo ""
                echo "$output"
                return 0
                ;;
        esac

        # Wait before next poll
        sleep $POLL_INTERVAL
        wait_time=$((wait_time + POLL_INTERVAL))
    done

    echo ""
    print_warning "Max wait time reached ($(format_duration $MAX_WAIT_TIME))"
    print_info "Job is still running, but test is ending"
}

get_best_params() {
    print_banner "Retrieving Best Parameters"

    print_step "Fetching best hyperparameters..."
    echo ""

    local output
    if ! output=$($TLI_BIN tune best --job-id "$JOB_ID" 2>&1); then
        print_error "Failed to get best parameters"
        echo "$output"
        exit 1
    fi

    echo "$output"
    echo ""

    # Optionally export to file
    local export_file="$PROJECT_ROOT/best_params_${JOB_ID:0:8}.yaml"
    print_step "Exporting best parameters to file..."

    if ! output=$($TLI_BIN tune best --job-id "$JOB_ID" --export "$export_file" 2>&1); then
        print_warning "Failed to export parameters"
    else
        print_success "Parameters exported to: $export_file"
    fi
}

stop_tuning_job() {
    local reason=${1:-"User requested stop"}

    print_banner "Stopping Tuning Job"

    print_step "Sending stop request..."
    echo ""

    local output
    if ! output=$($TLI_BIN tune stop --job-id "$JOB_ID" --reason "$reason" 2>&1); then
        print_error "Failed to stop job"
        echo "$output"
        return 1
    fi

    echo "$output"
    echo ""

    print_success "Job stopped successfully"
}

# ============================================================================
# Test Execution Options
# ============================================================================

run_full_workflow() {
    print_banner "Full TLI Tuning Workflow Test"

    check_prerequisites
    start_tuning_job
    poll_status
    get_best_params

    print_banner "Test Completed Successfully"
    print_success "Full workflow executed without errors"

    local total_time=$(elapsed_time)
    print_info "Total duration: $(format_duration $total_time)"
}

run_quick_test() {
    print_banner "Quick TLI Tuning Test (Start Only)"

    check_prerequisites
    start_tuning_job

    echo ""
    print_success "Job started successfully"
    print_info "Monitor with: $TLI_BIN tune status --job-id $JOB_ID"
    print_info "Get results: $TLI_BIN tune best --job-id $JOB_ID"
    print_info "Stop job: $TLI_BIN tune stop --job-id $JOB_ID"
}

show_help() {
    cat << EOF
Usage: $0 [OPTIONS]

Test the full TLI hyperparameter tuning workflow.

OPTIONS:
    --full          Run full workflow (start → poll → results)
    --quick         Quick test (start job only, no polling)
    --check-only    Check prerequisites only
    --model TYPE    Model type (default: DQN)
    --trials N      Number of trials (default: 5)
    --help          Show this help message

EXAMPLES:
    # Full workflow test (default)
    $0

    # Quick start test
    $0 --quick

    # Check prerequisites only
    $0 --check-only

    # Custom model and trials
    $0 --model PPO --trials 10

ENVIRONMENT:
    TLI_BIN         Path to TLI binary (default: target/release/tli)
    JWT_TOKEN_FILE  Path to JWT token (default: ~/.foxhunt/jwt_token)
    POLL_INTERVAL   Status polling interval in seconds (default: 5)
    MAX_WAIT_TIME   Maximum wait time in seconds (default: 300)

PREREQUISITES:
    1. TLI binary built (cargo build --release -p tli)
    2. JWT token exists (~/.foxhunt/jwt_token)
    3. API Gateway running (http://localhost:8080)
    4. ML Training Service running (http://localhost:8095)
    5. Test data available (optional)

EOF
}

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

main() {
    local mode="full"

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --full)
                mode="full"
                shift
                ;;
            --quick)
                mode="quick"
                shift
                ;;
            --check-only)
                mode="check"
                shift
                ;;
            --model)
                MODEL_TYPE="$2"
                shift 2
                ;;
            --trials)
                NUM_TRIALS="$2"
                shift 2
                ;;
            --help|-h)
                show_help
                exit 0
                ;;
            *)
                print_error "Unknown option: $1"
                echo ""
                show_help
                exit 1
                ;;
        esac
    done

    # Execute based on mode
    case $mode in
        full)
            run_full_workflow
            ;;
        quick)
            run_quick_test
            ;;
        check)
            check_prerequisites
            print_success "Prerequisites check complete"
            ;;
        *)
            print_error "Invalid mode: $mode"
            exit 1
            ;;
    esac
}

# Run main function
main "$@"