#!/bin/bash
#
# Databento Download Planning Tool
#
# Purpose: Estimate costs, check credit balance, and generate download plan
#          before actually downloading data from Databento.
#
# Usage:
#   ./scripts/plan_databento_download.sh \
#     --symbols "ES.FUT,NQ.FUT" \
#     --days 5 \
#     --schema ohlcv-1m \
#     [--start-date 2024-01-02]
#
# Features:
#   - Cost estimation based on historical data
#   - Credit balance check (80%, 90% alerts)
#   - Download command generation
#   - Budget safety checks
#

set -euo pipefail

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

COST_TRACKING_FILE="${COST_TRACKING_FILE:-./COST_TRACKING.md}"
TOTAL_CREDITS="${TOTAL_CREDITS:-125.00}"
WARN_THRESHOLD_PCT="${WARN_THRESHOLD_PCT:-80}"
CRITICAL_THRESHOLD_PCT="${CRITICAL_THRESHOLD_PCT:-90}"

# Schema cost estimates (per day, per symbol)
declare -A SCHEMA_COSTS=(
    ["ohlcv-1m"]="0.00002"
    ["ohlcv-1s"]="0.0004"
    ["tbbo"]="0.0001"
    ["trades"]="0.0012"
    ["mbp-1"]="0.05"
    ["mbp-10"]="1.0"
    ["mbo"]="3.0"
)

# Schema file size estimates (per day, per symbol, in KB)
declare -A SCHEMA_SIZES=(
    ["ohlcv-1m"]="90"
    ["ohlcv-1s"]="1400"
    ["tbbo"]="50"
    ["trades"]="240"
    ["mbp-1"]="4800"
    ["mbp-10"]="48000"
    ["mbo"]="100000"
)

# ============================================================================
# Color Codes
# ============================================================================

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

# ============================================================================
# Functions
# ============================================================================

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

Databento Download Planning Tool - Estimate costs before downloading

OPTIONS:
    --symbols SYMBOLS       Comma-separated list of symbols (e.g., "ES.FUT,NQ.FUT")
    --days DAYS             Number of days to download
    --schema SCHEMA         Schema type: ohlcv-1m, ohlcv-1s, tbbo, trades, mbp-1, mbp-10, mbo
    --start-date DATE       (Optional) Start date in YYYY-MM-DD format
    --parent                (Optional) Use parent symbol type (WARNING: increases cost 10-60x)
    --help                  Show this help message

EXAMPLES:
    # Single symbol, 5 days, OHLCV-1m
    $0 --symbols "ES.FUT" --days 5 --schema ohlcv-1m

    # Multiple symbols, 10 days, trades
    $0 --symbols "ES.FUT,NQ.FUT,RTY.FUT" --days 10 --schema trades

    # With specific start date
    $0 --symbols "ES.FUT" --days 7 --schema ohlcv-1m --start-date 2024-01-02

SUPPORTED SCHEMAS:
    ohlcv-1m    1-minute OHLCV bars (CHEAPEST, ~\$0.00002/day)
    ohlcv-1s    1-second OHLCV bars (~\$0.0004/day)
    tbbo        Top of book quotes (~\$0.0001/day)
    trades      All trades (~\$0.0012/day)
    mbp-1       Level 2, 1 level (~\$0.05/day)
    mbp-10      Level 2, 10 levels (~\$1.00/day)
    mbo         Level 3, full depth (MOST EXPENSIVE, ~\$3.00/day)

EOF
    exit 1
}

log_info() {
    echo -e "${BLUE}[INFO]${NC} $*"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $*"
}

log_warn() {
    echo -e "${YELLOW}[WARNING]${NC} $*"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $*"
}

# Calculate credits used from COST_TRACKING.md
calculate_credits_used() {
    if [ ! -f "$COST_TRACKING_FILE" ]; then
        log_warn "Cost tracking file not found: $COST_TRACKING_FILE"
        echo "0.00"
        return
    fi

    # Extract "Estimated Cost: $X.XX - $Y.YY" lines and sum the high estimates
    local total
    total=$(grep -oP 'Estimated Cost: \$[\d.]+ - \$\K[\d.]+' "$COST_TRACKING_FILE" 2>/dev/null | \
            awk '{sum+=$1} END {printf "%.4f", sum}')

    # If no matches, return 0
    if [ -z "$total" ]; then
        echo "0.00"
    else
        echo "$total"
    fi
}

# Check if credits are available for this download
check_credit_budget() {
    local estimated_cost=$1
    local credits_used
    local credits_remaining
    local usage_pct

    credits_used=$(calculate_credits_used)
    credits_remaining=$(echo "$TOTAL_CREDITS - $credits_used" | bc -l)
    usage_pct=$(echo "scale=2; ($credits_used / $TOTAL_CREDITS) * 100" | bc -l)

    echo ""
    echo "======================================================================"
    echo "                     CREDIT BUDGET STATUS"
    echo "======================================================================"
    printf "%-30s: \$%.4f\n" "Total Credits" "$TOTAL_CREDITS"
    printf "%-30s: \$%.4f (%.1f%%)\n" "Credits Used" "$credits_used" "$usage_pct"
    printf "%-30s: \$%.4f (%.1f%%)\n" "Credits Remaining" "$credits_remaining" "$(echo "100 - $usage_pct" | bc -l)"
    printf "%-30s: \$%.4f\n" "Estimated Cost (This Download)" "$estimated_cost"
    echo "======================================================================"

    # Check alert thresholds
    if (( $(echo "$usage_pct >= $CRITICAL_THRESHOLD_PCT" | bc -l) )); then
        log_error "🚨 RED ALERT: Credit usage at ${usage_pct}% (>${CRITICAL_THRESHOLD_PCT}%)"
        log_error "Emergency halt recommended. Executive approval required."
        return 2
    elif (( $(echo "$usage_pct >= $WARN_THRESHOLD_PCT" | bc -l) )); then
        log_warn "🟠 ORANGE ALERT: Credit usage at ${usage_pct}% (>${WARN_THRESHOLD_PCT}%)"
        log_warn "Review download necessity before proceeding."
        return 1
    else
        log_success "✅ GREEN: Credit usage at ${usage_pct}% (budget healthy)"
        return 0
    fi
}

# Estimate cost for download
estimate_cost() {
    local symbols=$1
    local days=$2
    local schema=$3
    local use_parent=$4

    # Count symbols
    IFS=',' read -ra symbol_array <<< "$symbols"
    local num_symbols=${#symbol_array[@]}

    # Get base cost per day per symbol
    local base_cost_per_day=${SCHEMA_COSTS[$schema]}

    # Apply parent symbol multiplier if enabled
    local parent_multiplier="1"
    if [ "$use_parent" = true ]; then
        log_warn "Parent symbol type selected: Cost multiplier 15x applied" >&2
        parent_multiplier="15"
    fi

    # Calculate total cost
    local total_cost
    total_cost=$(echo "scale=4; $num_symbols * $days * $base_cost_per_day * $parent_multiplier" | bc -l)

    # Estimate file size (in MB)
    local base_size_kb=${SCHEMA_SIZES[$schema]}
    local total_size_kb
    total_size_kb=$(echo "scale=2; $num_symbols * $days * $base_size_kb * $parent_multiplier" | bc -l)
    local total_size_mb
    total_size_mb=$(echo "scale=2; $total_size_kb / 1024" | bc -l)

    echo "$total_cost $total_size_mb"
}

# Generate download commands
generate_download_commands() {
    local symbols=$1
    local days=$2
    local schema=$3
    local start_date=$4
    local use_parent=$5

    # Parse symbols
    IFS=',' read -ra symbol_array <<< "$symbols"

    # Determine symbol type
    local stype_in="continuous"
    if [ "$use_parent" = true ]; then
        stype_in="parent"
    fi

    echo ""
    echo "======================================================================"
    echo "                   DOWNLOAD COMMANDS"
    echo "======================================================================"
    echo ""

    for symbol in "${symbol_array[@]}"; do
        # Trim whitespace
        symbol=$(echo "$symbol" | xargs)

        for ((day=0; day<days; day++)); do
            # Calculate dates
            if [ -n "$start_date" ]; then
                local current_date
                current_date=$(date -d "$start_date + $day days" +%Y-%m-%d)
            else
                local current_date
                current_date=$(date -d "2024-01-02 + $day days" +%Y-%m-%d)
            fi

            local next_date
            next_date=$(date -d "$current_date + 1 day" +%Y-%m-%d)

            # Generate filename
            local output_file="test_data/real/databento/${symbol}_${schema}_${current_date}.dbn"

            # Generate curl command
            cat << EOF
# Download: $symbol, $current_date ($schema)
curl -u "\${DATABENTO_API_KEY}:" \\
  "https://hist.databento.com/v0/timeseries.get_range?\\
dataset=GLBX.MDP3&\\
symbols=${symbol}&\\
schema=${schema}&\\
start=${current_date}T00:00:00Z&\\
end=${next_date}T00:00:00Z&\\
encoding=dbn&\\
stype_in=${stype_in}" \\
  --output "$output_file"

# Verify download
ls -lh "$output_file"

# Update COST_TRACKING.md
echo "## Download $(date +%Y-%m-%d): $symbol ($schema)" >> COST_TRACKING.md
echo "- Symbol: $symbol" >> COST_TRACKING.md
echo "- Schema: $schema" >> COST_TRACKING.md
echo "- Date: $current_date" >> COST_TRACKING.md
echo "- File: $output_file" >> COST_TRACKING.md
echo "" >> COST_TRACKING.md

EOF
        done
    done

    echo "======================================================================"
}

# ============================================================================
# Main
# ============================================================================

main() {
    local symbols=""
    local days=""
    local schema=""
    local start_date=""
    local use_parent=false

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --symbols)
                symbols="$2"
                shift 2
                ;;
            --days)
                days="$2"
                shift 2
                ;;
            --schema)
                schema="$2"
                shift 2
                ;;
            --start-date)
                start_date="$2"
                shift 2
                ;;
            --parent)
                use_parent=true
                shift
                ;;
            --help)
                usage
                ;;
            *)
                log_error "Unknown option: $1"
                usage
                ;;
        esac
    done

    # Validate required arguments
    if [ -z "$symbols" ] || [ -z "$days" ] || [ -z "$schema" ]; then
        log_error "Missing required arguments"
        usage
    fi

    # Validate schema
    if [ -z "${SCHEMA_COSTS[$schema]:-}" ]; then
        log_error "Invalid schema: $schema"
        log_error "Supported schemas: ${!SCHEMA_COSTS[*]}"
        exit 1
    fi

    # Display plan header
    echo ""
    echo "======================================================================"
    echo "              DATABENTO DOWNLOAD PLAN"
    echo "======================================================================"
    echo "Symbols:      $symbols"
    echo "Days:         $days"
    echo "Schema:       $schema"
    echo "Start Date:   ${start_date:-2024-01-02 (default)}"
    echo "Symbol Type:  $([ "$use_parent" = true ] && echo "parent (⚠️ 15x cost!)" || echo "continuous")"
    echo "======================================================================"

    # Estimate cost
    log_info "Estimating cost..."
    read -r estimated_cost estimated_size_mb <<< "$(estimate_cost "$symbols" "$days" "$schema" "$use_parent")"

    echo ""
    echo "======================================================================"
    echo "                     COST ESTIMATE"
    echo "======================================================================"
    printf "%-30s: \$%.4f\n" "Estimated Cost (High)" "$estimated_cost"
    printf "%-30s: %.2f MB\n" "Estimated File Size" "$estimated_size_mb"
    echo "======================================================================"

    # Check budget
    if ! check_credit_budget "$estimated_cost"; then
        local exit_code=$?
        if [ $exit_code -eq 2 ]; then
            log_error "Download REJECTED due to credit exhaustion risk"
            exit 1
        else
            log_warn "Download requires review before proceeding"
        fi
    fi

    # Generate download commands
    generate_download_commands "$symbols" "$days" "$schema" "$start_date" "$use_parent"

    # Final recommendation
    echo ""
    echo "======================================================================"
    echo "                     RECOMMENDATION"
    echo "======================================================================"

    if (( $(echo "$estimated_cost < 0.01" | bc -l) )); then
        log_success "✅ APPROVE: Cost <\$0.01, proceed immediately"
    elif (( $(echo "$estimated_cost < 0.10" | bc -l) )); then
        log_info "📋 REVIEW: Cost \$0.01-\$0.10, check budget and prioritize"
    elif (( $(echo "$estimated_cost < 1.00" | bc -l) )); then
        log_warn "🟡 REQUIRE APPROVAL: Cost \$0.10-\$1.00, justify and review"
    else
        log_error "🚨 EXECUTIVE APPROVAL: Cost >\$1.00, high-cost download"
    fi

    echo ""
    echo "Next Steps:"
    echo "1. Review cost estimate and budget status"
    echo "2. If approved, copy and execute download commands above"
    echo "3. Verify file sizes match estimates"
    echo "4. Update COST_TRACKING.md with actual costs"
    echo "======================================================================"
    echo ""
}

# Run main function
main "$@"