#!/bin/bash
# Foxhunt HFT Trading System - Database Migration Script

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
DEPLOY_DIR="/opt/foxhunt"
LOG_FILE="/home/jgrusewski/Work/foxhunt/logs-migration.log"

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

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

error() {
    log "${RED}ERROR: $1${NC}"
    exit 1
}

warning() {
    log "${YELLOW}WARNING: $1${NC}"
}

info() {
    log "${BLUE}INFO: $1${NC}"
}

success() {
    log "${GREEN}SUCCESS: $1${NC}"
}

# Load environment variables
load_env() {
    if [[ -f "${DEPLOY_DIR}/docker/.env" ]]; then
        source "${DEPLOY_DIR}/docker/.env"
    elif [[ -f "${PROJECT_ROOT}/docker/.env" ]]; then
        source "${PROJECT_ROOT}/docker/.env"
    else
        error "Environment file not found"
    fi
}

# Wait for database to be ready
wait_for_database() {
    info "Waiting for PostgreSQL to be ready..."

    local max_attempts=30
    local attempt=0

    while [[ $attempt -lt $max_attempts ]]; do
        if docker exec foxhunt-postgres pg_isready -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" &>/dev/null; then
            success "PostgreSQL is ready"
            return 0
        fi

        sleep 2
        ((attempt++))
        info "Waiting for PostgreSQL... (${attempt}/${max_attempts})"
    done

    error "PostgreSQL failed to become ready within expected time"
}

# Create database backup
create_backup() {
    info "Creating database backup..."

    local backup_file="/opt/foxhunt/backups/db-backup-$(date +%Y%m%d_%H%M%S).sql"
    mkdir -p "$(dirname "$backup_file")"

    docker exec foxhunt-postgres pg_dump -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" > "$backup_file"

    if [[ -f "$backup_file" && -s "$backup_file" ]]; then
        success "Database backup created: $backup_file"
        echo "$backup_file"
    else
        error "Failed to create database backup"
    fi
}

# Run PostgreSQL migrations
run_postgres_migrations() {
    info "Running PostgreSQL migrations..."

    local migrations_dir="${PROJECT_ROOT}/migrations"

    if [[ ! -d "$migrations_dir" ]]; then
        warning "No migrations directory found at $migrations_dir"
        return 0
    fi

    # Create migrations table if it doesn't exist
    docker exec foxhunt-postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -c "
        CREATE TABLE IF NOT EXISTS schema_migrations (
            version VARCHAR(255) PRIMARY KEY,
            applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
    "

    # Run migrations in order
    for migration_file in "$migrations_dir"/*.sql; do
        if [[ -f "$migration_file" ]]; then
            local version=$(basename "$migration_file" .sql)

            # Check if migration already applied
            local applied=$(docker exec foxhunt-postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -t -c "
                SELECT COUNT(*) FROM schema_migrations WHERE version = '$version';
            " | xargs)

            if [[ "$applied" == "0" ]]; then
                info "Applying migration: $version"

                # Run migration
                docker exec -i foxhunt-postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" < "$migration_file"

                # Record migration
                docker exec foxhunt-postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -c "
                    INSERT INTO schema_migrations (version) VALUES ('$version');
                "

                success "Migration applied: $version"
            else
                info "Migration already applied: $version"
            fi
        fi
    done

    success "PostgreSQL migrations completed"
}

# Setup InfluxDB
setup_influxdb() {
    info "Setting up InfluxDB..."

    # Wait for InfluxDB to be ready
    local max_attempts=30
    local attempt=0

    while [[ $attempt -lt $max_attempts ]]; do
        if curl -s http://localhost:8086/ping &>/dev/null; then
            success "InfluxDB is ready"
            break
        fi

        sleep 2
        ((attempt++))
        info "Waiting for InfluxDB... (${attempt}/${max_attempts})"
    done

    if [[ $attempt -eq $max_attempts ]]; then
        error "InfluxDB failed to become ready"
    fi

    # Create buckets if they don't exist
    local buckets=("market_data" "trading_metrics" "system_metrics" "latency_metrics")

    for bucket in "${buckets[@]}"; do
        info "Creating InfluxDB bucket: $bucket"

        # Use InfluxDB API to create bucket
        curl -s -X POST "http://localhost:8086/api/v2/buckets" \
            -H "Authorization: Token ${INFLUXDB_TOKEN}" \
            -H "Content-Type: application/json" \
            -d "{
                \"name\": \"$bucket\",
                \"orgID\": \"$(curl -s "http://localhost:8086/api/v2/orgs?org=${INFLUXDB_ORG}" -H "Authorization: Token ${INFLUXDB_TOKEN}" | grep -o '"id":"[^"]*"' | head -1 | cut -d'"' -f4)\",
                \"retentionRules\": [{\"everySeconds\": 2592000}]
            }" || warning "Bucket $bucket may already exist"
    done

    success "InfluxDB setup completed"
}

# Setup ClickHouse
setup_clickhouse() {
    info "Setting up ClickHouse..."

    # Wait for ClickHouse to be ready
    local max_attempts=30
    local attempt=0

    while [[ $attempt -lt $max_attempts ]]; do
        if curl -s http://localhost:8123/ping &>/dev/null; then
            success "ClickHouse is ready"
            break
        fi

        sleep 2
        ((attempt++))
        info "Waiting for ClickHouse... (${attempt}/${max_attempts})"
    done

    if [[ $attempt -eq $max_attempts ]]; then
        error "ClickHouse failed to become ready"
    fi

    # Create database and tables
    info "Creating ClickHouse database and tables..."

    # Create database
    curl -s -X POST "http://localhost:8123/" -d "CREATE DATABASE IF NOT EXISTS ${CLICKHOUSE_DB}"

    # Create market data table
    curl -s -X POST "http://localhost:8123/" -d "
        CREATE TABLE IF NOT EXISTS ${CLICKHOUSE_DB}.market_data (
            timestamp DateTime64(3),
            symbol String,
            price Float64,
            volume UInt64,
            side Enum8('buy' = 1, 'sell' = 2),
            exchange String
        ) ENGINE = MergeTree()
        PARTITION BY toYYYYMM(timestamp)
        ORDER BY (symbol, timestamp)
    "

    # Create trades table
    curl -s -X POST "http://localhost:8123/" -d "
        CREATE TABLE IF NOT EXISTS ${CLICKHOUSE_DB}.trades (
            timestamp DateTime64(3),
            trade_id String,
            symbol String,
            price Float64,
            quantity Float64,
            side Enum8('buy' = 1, 'sell' = 2),
            strategy String,
            pnl Float64
        ) ENGINE = MergeTree()
        PARTITION BY toYYYYMM(timestamp)
        ORDER BY (symbol, timestamp)
    "

    success "ClickHouse setup completed"
}

# Show migration summary
show_summary() {
    local backup_file="$1"

    info "Migration Summary"
    echo "===================="
    echo "Migration completed at: $(date)"
    echo "Database backup: $backup_file"
    echo ""
    echo "Databases:"
    echo "  - PostgreSQL: Ready for application data"
    echo "  - InfluxDB: Ready for time-series data"
    echo "  - ClickHouse: Ready for analytics"
    echo "  - Redis: Ready for caching"
    echo ""
    echo "To verify:"
    echo "  - PostgreSQL: docker exec foxhunt-postgres psql -U ${POSTGRES_USER} -d ${POSTGRES_DB} -c '\\dt'"
    echo "  - InfluxDB: curl http://localhost:8086/ping"
    echo "  - ClickHouse: curl http://localhost:8123/ping"
    echo "  - Redis: docker exec foxhunt-redis redis-cli ping"
}

# Main migration function
main() {
    info "Starting database migration..."

    # Check if running as root
    if [[ $EUID -ne 0 ]]; then
        error "This script must be run as root (use sudo)"
    fi

    load_env
    wait_for_database

    local backup_file
    backup_file=$(create_backup)

    run_postgres_migrations
    setup_influxdb
    setup_clickhouse

    success "Database migration completed successfully!"
    show_summary "$backup_file"
}

# Handle script arguments
case "${1:-migrate}" in
    "migrate")
        main
        ;;
    "backup")
        load_env
        wait_for_database
        create_backup
        ;;
    "postgres")
        load_env
        wait_for_database
        run_postgres_migrations
        ;;
    "influxdb")
        load_env
        setup_influxdb
        ;;
    "clickhouse")
        load_env
        setup_clickhouse
        ;;
    *)
        echo "Usage: $0 {migrate|backup|postgres|influxdb|clickhouse}"
        exit 1
        ;;
esac