#!/bin/bash

# =============================================================================
# UV Migration Script
# =============================================================================
# Migrates from pip/venv to uv package manager
# Includes cleanup, verification, and safety checks
# =============================================================================

set -e  # Exit on any error

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

# Script version
VERSION="1.0.0"

# Print header
echo -e "${BLUE}=================================${NC}"
echo -e "${BLUE}   UV Migration Script v${VERSION}${NC}"
echo -e "${BLUE}=================================${NC}"
echo ""

# Function to print colored output
print_success() {
    echo -e "${GREEN}✅ $1${NC}"
}

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

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

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

# Function to check if command exists
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# Function to get file size in human readable format
get_size() {
    if [[ "$OSTYPE" == "darwin"* ]]; then
        # macOS
        du -h "$1" 2>/dev/null | cut -f1 || echo "0B"
    else
        # Linux
        du -sh "$1" 2>/dev/null | cut -f1 || echo "0B"
    fi
}

# Function to count files
count_files() {
    find "$1" -type f 2>/dev/null | wc -l | tr -d ' '
}

# Check if we're in the project root
if [[ ! -f "pyproject.toml" ]]; then
    print_error "pyproject.toml not found. Please run this script from the project root."
    exit 1
fi

print_info "Starting migration process..."
echo ""

# =============================================================================
# Step 1: Pre-migration Analysis
# =============================================================================
echo -e "${BLUE}📊 Pre-migration Analysis${NC}"
echo "----------------------------"

# Count cache files before cleanup
PYCACHE_DIRS=$(find . -path "./.venv" -prune -o -name "__pycache__" -type d -print | wc -l | tr -d ' ')
PYC_FILES=$(find . -path "./.venv" -prune -o -name "*.pyc" -type f -print | wc -l | tr -d ' ')

echo "📁 Cache directories found: $PYCACHE_DIRS"
echo "📄 .pyc files found: $PYC_FILES"

# Check if pytest cache exists
if [[ -d ".pytest_cache" ]]; then
    PYTEST_SIZE=$(get_size ".pytest_cache")
    echo "🧪 Pytest cache size: $PYTEST_SIZE"
else
    echo "🧪 No pytest cache found"
fi

# Check for old venv
if [[ -d "venv" ]]; then
    print_warning "Old 'venv' directory still exists!"
    VENV_SIZE=$(get_size "venv")
    echo "📦 Old venv size: $VENV_SIZE"
else
    print_success "No old 'venv' directory found"
fi

# Check for requirements.txt
if [[ -f "requirements.txt" ]]; then
    print_warning "requirements.txt still exists (should be removed after migration)"
else
    print_success "requirements.txt already removed"
fi

echo ""

# =============================================================================
# Step 2: UV Installation Check
# =============================================================================
echo -e "${BLUE}🔧 UV Installation Check${NC}"
echo "----------------------------"

if command_exists uv; then
    UV_VERSION=$(uv --version | head -1)
    print_success "UV installed: $UV_VERSION"
else
    print_error "UV not found! Installing UV..."
    if command_exists curl; then
        curl -LsSf https://astral.sh/uv/install.sh | sh
        # Add to PATH for this session
        export PATH="$HOME/.local/bin:$PATH"
        if command_exists uv; then
            UV_VERSION=$(uv --version | head -1)
            print_success "UV installed successfully: $UV_VERSION"
        else
            print_error "UV installation failed!"
            exit 1
        fi
    else
        print_error "curl not found! Please install UV manually."
        exit 1
    fi
fi

echo ""

# =============================================================================
# Step 3: Cache Cleanup
# =============================================================================
echo -e "${BLUE}🧹 Cache Cleanup${NC}"
echo "-------------------"

# Ask for confirmation unless --force flag is provided
if [[ "$1" != "--force" ]]; then
    echo "This will delete:"
    echo "  - $PYCACHE_DIRS __pycache__ directories"
    echo "  - $PYC_FILES .pyc files"
    echo "  - .pytest_cache directory (if exists)"
    echo ""
    read -p "Continue with cleanup? (y/N): " -n 1 -r
    echo
    if [[ ! $REPLY =~ ^[Yy]$ ]]; then
        print_warning "Cleanup cancelled by user"
        exit 0
    fi
fi

print_info "Starting cleanup..."

# Remove __pycache__ directories (excluding .venv)
DELETED_DIRS=0
while IFS= read -r -d '' dir; do
    if [[ "$dir" != *"/.venv/"* ]]; then
        rm -rf "$dir"
        ((DELETED_DIRS++))
    fi
done < <(find . -name "__pycache__" -type d -print0)

# Remove .pyc files (excluding .venv)
DELETED_FILES=0
while IFS= read -r -d '' file; do
    if [[ "$file" != *"/.venv/"* ]]; then
        rm -f "$file"
        ((DELETED_FILES++))
    fi
done < <(find . -name "*.pyc" -type f -print0)

# Remove pytest cache
if [[ -d ".pytest_cache" ]]; then
    rm -rf ".pytest_cache"
    print_success "Removed .pytest_cache directory"
fi

# Remove .pyo files if any
find . -path "./.venv" -prune -o -name "*.pyo" -type f -delete

print_success "Cleanup completed!"
echo "  - Removed $DELETED_DIRS __pycache__ directories"
echo "  - Removed $DELETED_FILES .pyc files"

echo ""

# =============================================================================
# Step 4: UV Environment Verification
# =============================================================================
echo -e "${BLUE}🔍 UV Environment Verification${NC}"
echo "--------------------------------"

# Check if .venv exists
if [[ -d ".venv" ]]; then
    print_success "UV environment (.venv) exists"
else
    print_warning "UV environment not found. Creating..."
    uv sync
fi

# Check if uv.lock exists
if [[ -f "uv.lock" ]]; then
    print_success "uv.lock file exists"
else
    print_warning "uv.lock not found. Running uv sync..."
    uv sync
fi

# Test imports
print_info "Testing key imports..."
if uv run python -c "import hyperliquid, pandas, requests; print('All imports successful!')" >/dev/null 2>&1; then
    print_success "All key dependencies working"
else
    print_error "Import test failed!"
    exit 1
fi

echo ""

# =============================================================================
# Step 5: Old Files Cleanup (Optional)
# =============================================================================
echo -e "${BLUE}🗑️  Old Files Cleanup${NC}"
echo "-----------------------"

# Check for old venv directory
if [[ -d "venv" ]]; then
    VENV_SIZE=$(get_size "venv")
    print_warning "Old 'venv' directory found (${VENV_SIZE})"
    
    if [[ "$1" != "--force" ]]; then
        read -p "Delete old 'venv' directory? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            rm -rf "venv"
            print_success "Removed old 'venv' directory"
        else
            print_warning "Keeping old 'venv' directory"
        fi
    else
        rm -rf "venv"
        print_success "Removed old 'venv' directory"
    fi
else
    print_success "No old 'venv' directory found"
fi

# Check for requirements.txt
if [[ -f "requirements.txt" ]]; then
    print_warning "requirements.txt still exists"
    
    if [[ "$1" != "--force" ]]; then
        read -p "Delete requirements.txt? (y/N): " -n 1 -r
        echo
        if [[ $REPLY =~ ^[Yy]$ ]]; then
            rm -f "requirements.txt"
            print_success "Removed requirements.txt"
        else
            print_warning "Keeping requirements.txt"
        fi
    else
        rm -f "requirements.txt"
        print_success "Removed requirements.txt"
    fi
else
    print_success "requirements.txt already removed"
fi

echo ""

# =============================================================================
# Step 6: Final Verification
# =============================================================================
echo -e "${BLUE}✅ Final Verification${NC}"
echo "-----------------------"

# Count remaining cache files
REMAINING_PYCACHE=$(find . -path "./.venv" -prune -o -name "__pycache__" -type d -print | wc -l | tr -d ' ')
REMAINING_PYC=$(find . -path "./.venv" -prune -o -name "*.pyc" -type f -print | wc -l | tr -d ' ')

if [[ $REMAINING_PYCACHE -eq 0 && $REMAINING_PYC -eq 0 ]]; then
    print_success "All cache files cleaned"
else
    print_warning "$REMAINING_PYCACHE cache dirs and $REMAINING_PYC .pyc files remaining"
fi

# Test UV commands
print_info "Testing UV commands..."

# Test uv run
if uv run python --version >/dev/null 2>&1; then
    print_success "uv run working"
else
    print_error "uv run failed"
fi

# Test project dependencies
if uv run python -c "import src; print('Project imports working')" >/dev/null 2>&1; then
    print_success "Project imports working"
else
    print_warning "Project imports need attention"
fi

echo ""

# =============================================================================
# Migration Summary
# =============================================================================
echo -e "${GREEN}🎉 Migration Summary${NC}"
echo "====================="
echo ""
print_success "UV migration completed successfully!"
echo ""
echo "📋 What was done:"
echo "  ✅ Cleaned up $DELETED_DIRS cache directories"
echo "  ✅ Removed $DELETED_FILES .pyc files"
echo "  ✅ Verified UV installation and environment"
echo "  ✅ Tested key dependencies"
echo ""
echo "🚀 Next steps:"
echo "  • Use 'uv run python script.py' instead of 'python script.py'"
echo "  • Use 'uv sync' to install/update dependencies"
echo "  • Use 'uv add package_name' to add new packages"
echo ""
echo "📚 Quick reference:"
echo "  uv sync                    # Install dependencies"
echo "  uv run python trading_bot.py  # Run your bot"
echo "  uv add requests            # Add new package"
echo "  uv remove requests         # Remove package"
echo ""
print_success "Happy coding with UV! 🚀"