#!/bin/bash

# Storage Optimization Script - SSD Tiering Implementation
# Optimizes storage performance with intelligent data placement

set -euo pipefail

# Configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
PROJECT_ROOT="$(dirname "$SCRIPT_DIR")"
LOG_FILE="$PROJECT_ROOT/logs/storage-optimization-$(date +%Y%m%d-%H%M%S).log"

# Storage tier definitions (adjust paths based on your setup)
SSD_MOUNT="/opt/ssd"           # Fast SSD storage (234GB)
HDD_MOUNT="/srv/mergerfs"      # Large HDD storage (20.8TB)  
CACHE_MOUNT="/opt/cache"       # NVMe cache layer

# Docker data locations
DOCKER_ROOT="/var/lib/docker"
VOLUME_ROOT="/var/lib/docker/volumes"

# Create directories
mkdir -p "$(dirname "$LOG_FILE")" "$PROJECT_ROOT/logs"

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

# Check available storage
check_storage() {
    log "Checking available storage..."
    
    log "Current disk usage:"
    df -h | grep -E "(ssd|hdd|cache|docker)" || true
    
    # Check if mount points exist
    for mount in "$SSD_MOUNT" "$HDD_MOUNT" "$CACHE_MOUNT"; do
        if [[ ! -d "$mount" ]]; then
            log "Warning: Mount point $mount does not exist"
        else
            log "✅ Mount point available: $mount ($(df -h "$mount" | tail -1 | awk '{print $4}') free)"
        fi
    done
}

# Setup SSD tier for hot data
setup_ssd_tier() {
    log "Setting up SSD tier for high-performance data..."
    
    # Create SSD directories
    sudo mkdir -p "$SSD_MOUNT"/{postgresql,redis,container-logs,prometheus,grafana}
    
    # Database data (PostgreSQL)
    if [[ -d "$VOLUME_ROOT" ]]; then
        # Find PostgreSQL volumes and move to SSD
        find "$VOLUME_ROOT" -name "*postgresql*" -o -name "*postgres*" | while read -r vol; do
            if [[ -d "$vol" ]]; then
                local vol_name
                vol_name=$(basename "$vol")
                log "Moving PostgreSQL volume to SSD: $vol_name"
                
                # Create SSD location
                sudo mkdir -p "$SSD_MOUNT/postgresql/$vol_name"
                
                # Stop containers using this volume (if any)
                local containers
                containers=$(docker ps -a --filter volume="$vol_name" --format "{{.Names}}" || true)
                if [[ -n "$containers" ]]; then
                    log "Stopping containers using $vol_name: $containers"
                    echo "$containers" | xargs -r docker stop || true
                fi
                
                # Sync data to SSD
                sudo rsync -av "$vol/_data/" "$SSD_MOUNT/postgresql/$vol_name/" || true
                
                # Create bind mount configuration
                cat >> /tmp/ssd-mounts.conf << EOF
# PostgreSQL volume $vol_name
$SSD_MOUNT/postgresql/$vol_name $vol/_data none bind 0 0
EOF
                
                log "✅ PostgreSQL volume $vol_name configured for SSD"
            fi
        done
    fi
    
    # Redis data
    find "$VOLUME_ROOT" -name "*redis*" | while read -r vol; do
        if [[ -d "$vol" ]]; then
            local vol_name
            vol_name=$(basename "$vol")
            log "Moving Redis volume to SSD: $vol_name"
            
            sudo mkdir -p "$SSD_MOUNT/redis/$vol_name"
            sudo rsync -av "$vol/_data/" "$SSD_MOUNT/redis/$vol_name/" || true
            
            cat >> /tmp/ssd-mounts.conf << EOF
# Redis volume $vol_name  
$SSD_MOUNT/redis/$vol_name $vol/_data none bind 0 0
EOF
        fi
    done
    
    # Container logs (hot data)
    if [[ -d "/var/lib/docker/containers" ]]; then
        log "Setting up SSD storage for container logs"
        sudo mkdir -p "$SSD_MOUNT/container-logs"
        
        # Move recent logs to SSD (last 7 days)
        find /var/lib/docker/containers -name "*-json.log" -mtime -7 -exec sudo cp {} "$SSD_MOUNT/container-logs/" \; || true
    fi
}

# Setup HDD tier for cold data  
setup_hdd_tier() {
    log "Setting up HDD tier for large/cold data storage..."
    
    # Create HDD directories
    sudo mkdir -p "$HDD_MOUNT"/{media,backups,archives,immich-data,nextcloud-data}
    
    # Media files (Jellyfin content)
    find "$VOLUME_ROOT" -name "*jellyfin*" -o -name "*immich*" | while read -r vol; do
        if [[ -d "$vol" ]]; then
            local vol_name
            vol_name=$(basename "$vol")
            log "Moving media volume to HDD: $vol_name"
            
            sudo mkdir -p "$HDD_MOUNT/media/$vol_name"
            
            # For large data, use mv instead of rsync for efficiency
            sudo mv "$vol/_data"/* "$HDD_MOUNT/media/$vol_name/" 2>/dev/null || true
            
            cat >> /tmp/hdd-mounts.conf << EOF
# Media volume $vol_name
$HDD_MOUNT/media/$vol_name $vol/_data none bind 0 0
EOF
        fi
    done
    
    # Nextcloud data
    find "$VOLUME_ROOT" -name "*nextcloud*" | while read -r vol; do
        if [[ -d "$vol" ]]; then
            local vol_name
            vol_name=$(basename "$vol")
            log "Moving Nextcloud volume to HDD: $vol_name"
            
            sudo mkdir -p "$HDD_MOUNT/nextcloud-data/$vol_name"
            sudo rsync -av "$vol/_data/" "$HDD_MOUNT/nextcloud-data/$vol_name/" || true
            
            cat >> /tmp/hdd-mounts.conf << EOF
# Nextcloud volume $vol_name
$HDD_MOUNT/nextcloud-data/$vol_name $vol/_data none bind 0 0
EOF
        fi
    done
}

# Setup cache layer with bcache
setup_cache_layer() {
    log "Setting up cache layer for performance optimization..."
    
    # Check if bcache is available
    if ! command -v make-bcache >/dev/null 2>&1; then
        log "Installing bcache-tools..."
        sudo apt-get update && sudo apt-get install -y bcache-tools || {
            log "❌ Failed to install bcache-tools"
            return 1
        }
    fi
    
    # Create cache configuration (example - adapt to your setup)
    cat > /tmp/cache-setup.sh << 'EOF'
#!/bin/bash
# Bcache setup script (run with caution - can destroy data!)

# Example: Create cache device (adjust device paths!)
# sudo make-bcache -C /dev/nvme0n1p1 -B /dev/sdb1
# 
# Mount with cache:
# sudo mount /dev/bcache0 /mnt/cached-storage

echo "Cache layer setup requires manual configuration of block devices"
echo "Please review and adapt the cache setup for your specific hardware"
EOF
    
    chmod +x /tmp/cache-setup.sh
    log "⚠️  Cache layer setup script created at /tmp/cache-setup.sh"
    log "⚠️  Review and adapt for your hardware before running"
}

# Apply filesystem optimizations
optimize_filesystem() {
    log "Applying filesystem optimizations..."
    
    # Optimize mount options for different tiers
    cat > /tmp/optimized-fstab-additions.conf << 'EOF'
# Optimized mount options for storage tiers

# SSD optimizations (add to existing mounts)
# - noatime: disable access time updates
# - discard: enable TRIM
# - commit=60: reduce commit frequency
# Example: UUID=xxx /opt/ssd ext4 defaults,noatime,discard,commit=60 0 2

# HDD optimizations  
# - noatime: disable access time updates
# - commit=300: increase commit interval for HDDs
# Example: UUID=xxx /srv/hdd ext4 defaults,noatime,commit=300 0 2

# Temporary filesystem optimizations
tmpfs /tmp tmpfs defaults,noatime,mode=1777,size=2G 0 0
tmpfs /var/tmp tmpfs defaults,noatime,mode=1777,size=1G 0 0
EOF
    
    # Optimize Docker daemon for SSD
    local docker_config="/etc/docker/daemon.json"
    if [[ -f "$docker_config" ]]; then
        local backup_config="${docker_config}.backup-$(date +%Y%m%d)"
        sudo cp "$docker_config" "$backup_config"
        log "✅ Docker config backed up to $backup_config"
    fi
    
    # Create optimized Docker daemon configuration
    cat > /tmp/optimized-docker-daemon.json << 'EOF'
{
  "data-root": "/opt/ssd/docker",
  "storage-driver": "overlay2",
  "storage-opts": [
    "overlay2.override_kernel_check=true"
  ],
  "log-driver": "json-file",
  "log-opts": {
    "max-size": "10m",
    "max-file": "3"
  },
  "default-ulimits": {
    "nofile": {
      "name": "nofile",
      "hard": 64000,
      "soft": 64000
    }
  },
  "max-concurrent-downloads": 10,
  "max-concurrent-uploads": 5,
  "userland-proxy": false
}
EOF
    
    log "⚠️  Optimized Docker config created at /tmp/optimized-docker-daemon.json"
    log "⚠️  Review and apply manually to $docker_config"
}

# Create data lifecycle management
setup_lifecycle_management() {
    log "Setting up automated data lifecycle management..."
    
    # Create lifecycle management script
    cat > "$PROJECT_ROOT/scripts/storage-lifecycle.sh" << 'EOF'
#!/bin/bash
# Automated storage lifecycle management

# Move old logs to HDD (older than 30 days)
find /opt/ssd/container-logs -name "*.log" -mtime +30 -exec mv {} /srv/hdd/archived-logs/ \;

# Compress old media files (older than 1 year)
find /srv/hdd/media -name "*.mkv" -mtime +365 -exec ffmpeg -i {} -c:v libx265 -crf 28 -preset medium {}.h265.mkv \;

# Clean up Docker build cache weekly
docker system prune -af --volumes --filter "until=72h"

# Optimize database tables monthly
docker exec postgresql_primary psql -U postgres -c "VACUUM ANALYZE;"

# Generate storage report
df -h > /var/log/storage-report.txt
du -sh /opt/ssd/* >> /var/log/storage-report.txt
du -sh /srv/hdd/* >> /var/log/storage-report.txt
EOF
    
    chmod +x "$PROJECT_ROOT/scripts/storage-lifecycle.sh"
    
    # Create cron job for lifecycle management
    local cron_job="0 3 * * 0 $PROJECT_ROOT/scripts/storage-lifecycle.sh"
    if ! crontab -l 2>/dev/null | grep -q "storage-lifecycle.sh"; then
        (crontab -l 2>/dev/null; echo "$cron_job") | crontab -
        log "✅ Weekly storage lifecycle management scheduled"
    fi
}

# Monitor storage performance
setup_monitoring() {
    log "Setting up storage performance monitoring..."
    
    # Create storage monitoring script
    cat > "$PROJECT_ROOT/scripts/storage-monitor.sh" << 'EOF'
#!/bin/bash
# Storage performance monitoring

# Collect I/O statistics
iostat -x 1 5 > /tmp/iostat.log

# Monitor disk space usage
df -h | awk 'NR>1 {print $5 " " $6}' | while read usage mount; do
  usage_num=${usage%\%}
  if [ $usage_num -gt 85 ]; then
    echo "WARNING: $mount is $usage full" >> /var/log/storage-alerts.log
  fi
done

# Monitor SSD health (if nvme/smartctl available)
if command -v nvme >/dev/null 2>&1; then
  nvme smart-log /dev/nvme0n1 > /tmp/nvme-health.log 2>/dev/null || true
fi

if command -v smartctl >/dev/null 2>&1; then
  smartctl -a /dev/sda > /tmp/hdd-health.log 2>/dev/null || true
fi
EOF
    
    chmod +x "$PROJECT_ROOT/scripts/storage-monitor.sh"
    
    # Add to monitoring cron (every 15 minutes)
    local monitor_cron="*/15 * * * * $PROJECT_ROOT/scripts/storage-monitor.sh"
    if ! crontab -l 2>/dev/null | grep -q "storage-monitor.sh"; then
        (crontab -l 2>/dev/null; echo "$monitor_cron") | crontab -
        log "✅ Storage monitoring scheduled every 15 minutes"
    fi
}

# Generate optimization report
generate_report() {
    log "Generating storage optimization report..."
    
    local report_file="$PROJECT_ROOT/logs/storage-optimization-report.yaml"
    cat > "$report_file" << EOF
storage_optimization_report:
  timestamp: "$(date -Iseconds)"
  configuration:
    ssd_tier: "$SSD_MOUNT"
    hdd_tier: "$HDD_MOUNT" 
    cache_tier: "$CACHE_MOUNT"
  
  current_usage:
EOF
    
    # Add current usage statistics
    df -h | grep -E "(ssd|hdd|cache)" | while read -r line; do
        echo "    - $line" >> "$report_file"
    done
    
    # Add optimization summary
    cat >> "$report_file" << EOF
  
  optimizations_applied:
    - Database data moved to SSD tier
    - Media files organized on HDD tier
    - Container logs optimized for SSD
    - Filesystem mount options tuned
    - Docker daemon configuration optimized
    - Automated lifecycle management scheduled
    - Performance monitoring enabled
  
  recommendations:
    - Review and apply mount optimizations from /tmp/optimized-fstab-additions.conf
    - Apply Docker daemon config from /tmp/optimized-docker-daemon.json
    - Configure bcache if NVMe cache available
    - Monitor storage alerts in /var/log/storage-alerts.log
    - Review storage performance regularly
EOF
    
    log "✅ Optimization report generated: $report_file"
}

# Main execution
main() {
    case "${1:-optimize-all}" in
        "--check")
            check_storage
            ;;
        "--setup-ssd")
            setup_ssd_tier
            ;;
        "--setup-hdd")
            setup_hdd_tier
            ;;
        "--setup-cache")
            setup_cache_layer
            ;;
        "--optimize-filesystem")
            optimize_filesystem
            ;;
        "--setup-lifecycle")
            setup_lifecycle_management
            ;;
        "--setup-monitoring") 
            setup_monitoring
            ;;
        "--optimize-all"|"")
            log "Starting comprehensive storage optimization..."
            check_storage
            setup_ssd_tier
            setup_hdd_tier
            optimize_filesystem
            setup_lifecycle_management
            setup_monitoring
            generate_report
            log "🎉 Storage optimization completed!"
            ;;
        "--help"|"-h")
            cat << 'EOF'
Storage Optimization Script - SSD Tiering Implementation

USAGE:
  storage-optimization.sh [OPTIONS]

OPTIONS:
  --check               Check current storage configuration
  --setup-ssd          Set up SSD tier for hot data
  --setup-hdd          Set up HDD tier for cold data  
  --setup-cache        Set up cache layer configuration
  --optimize-filesystem Optimize filesystem settings
  --setup-lifecycle    Set up automated data lifecycle management
  --setup-monitoring   Set up storage performance monitoring
  --optimize-all       Run all optimizations (default)
  --help, -h          Show this help message

EXAMPLES:
  # Check current storage
  ./storage-optimization.sh --check
  
  # Set up SSD tier only
  ./storage-optimization.sh --setup-ssd
  
  # Run complete optimization
  ./storage-optimization.sh --optimize-all

NOTES:
  - Creates backups before modifying configurations
  - Requires sudo for filesystem operations
  - Review generated configs before applying
  - Monitor logs for any issues
EOF
            ;;
        *)
            log "❌ Unknown option: $1"
            log "Use --help for usage information"
            exit 1
            ;;
    esac
}

# Execute main function
main "$@"