#!/bin/bash
# ================================================
# Math Platform - Database Backup Script
# ================================================

set -e

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

# Configuration
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
BACKUP_DIR="$PROJECT_ROOT/docker/backups"
DB_CONTAINER="math-postgres"
DB_USER="mathuser"
DB_NAME="mathdb"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
BACKUP_FILE="$BACKUP_DIR/mathdb_backup_$TIMESTAMP.sql"
COMPRESSED_FILE="$BACKUP_FILE.gz"

# Retention settings
RETENTION_DAYS=30
RETENTION_COUNT=20

echo -e "${BLUE}============================================${NC}"
echo -e "${BLUE}Math Platform - Database Backup${NC}"
echo -e "${BLUE}============================================${NC}"

# Create backup directory
mkdir -p "$BACKUP_DIR"

# Function to check if container is running
check_container() {
    if ! docker ps | grep -q $DB_CONTAINER; then
        echo -e "${RED}Error: Database container $DB_CONTAINER is not running!${NC}"
        exit 1
    fi
    echo -e "${GREEN}Database container is running${NC}"
}

# Function to create backup
create_backup() {
    echo -e "${YELLOW}Creating database backup...${NC}"
    echo -e "Timestamp: $TIMESTAMP"
    echo -e "Destination: $COMPRESSED_FILE"

    if docker exec $DB_CONTAINER pg_dump -U $DB_USER $DB_NAME > "$BACKUP_FILE" 2>/dev/null; then
        if [ -f "$BACKUP_FILE" ] && [ -s "$BACKUP_FILE" ]; then
            # Compress backup
            gzip "$BACKUP_FILE"

            # Get file size
            FILE_SIZE=$(du -h "$COMPRESSED_FILE" | cut -f1)
            echo -e "${GREEN}Backup created successfully!${NC}"
            echo -e "${GREEN}File size: $FILE_SIZE${NC}"
            echo -e "${GREEN}Location: $COMPRESSED_FILE${NC}"

            # Create checksum
            sha256sum "$COMPRESSED_FILE" > "$COMPRESSED_FILE.sha256"
            echo -e "${GREEN}Checksum: ${COMPRESSED_FILE}.sha256${NC}"
        else
            echo -e "${RED}Error: Backup file is empty or was not created!${NC}"
            rm -f "$BACKUP_FILE"
            exit 1
        fi
    else
        echo -e "${RED}Error: Failed to create database backup!${NC}"
        exit 1
    fi
}

# Function to restore backup
restore_backup() {
    local backup_file=$1

    if [ -z "$backup_file" ]; then
        echo -e "${RED}Error: Please specify backup file to restore${NC}"
        echo "Usage: $0 --restore <backup_file>"
        exit 1
    fi

    if [ ! -f "$backup_file" ]; then
        echo -e "${RED}Error: Backup file not found: $backup_file${NC}"
        exit 1
    fi

    echo -e "${YELLOW}Restoring database from backup...${NC}"
    echo -e "Source: $backup_file"
    echo -e "${RED}WARNING: This will overwrite the current database!${NC}"

    read -p "Are you sure? (yes/no): " confirm
    if [ "$confirm" != "yes" ]; then
        echo -e "${YELLOW}Restore cancelled${NC}"
        exit 0
    fi

    # Decompress if needed
    local temp_file="$backup_file"
    if [[ "$backup_file" == *.gz ]]; then
        temp_file="/tmp/restore_$(date +%s).sql"
        gunzip -c "$backup_file" > "$temp_file"
    fi

    # Drop existing database and recreate
    docker exec -i $DB_CONTAINER psql -U $DB_USER -d postgres <<-EOF
        DROP DATABASE IF EXISTS $DB_NAME;
        CREATE DATABASE $DB_NAME;
EOF

    # Restore backup
    docker exec -i $DB_CONTAINER psql -U $DB_USER $DB_NAME < "$temp_file"

    # Cleanup
    if [ "$temp_file" != "$backup_file" ]; then
        rm -f "$temp_file"
    fi

    echo -e "${GREEN}Database restored successfully!${NC}"
}

# Function to list backups
list_backups() {
    echo -e "\n${BLUE}============================================${NC}"
    echo -e "${BLUE}Available Backups${NC}"
    echo -e "${BLUE}============================================${NC}"

    if [ "$(ls -A $BACKUP_DIR 2>/dev/null)" ]; then
        printf "%-40s %-15s %-10s\n" "File Name" "Date" "Size"
        printf "%-40s %-15s %-10s\n" "---------" "----" "----"

        for file in "$BACKUP_DIR"/mathdb_backup_*.sql.gz; do
            if [ -f "$file" ]; then
                filename=$(basename "$file")
                date=$(echo $filename | grep -oP '\d{8}_\d{6}' | sed 's/_/ /')
                size=$(du -h "$file" | cut -f1)
                printf "%-40s %-15s %-10s\n" "$filename" "$date" "$size"
            fi
        done
    else
        echo -e "${YELLOW}No backups found${NC}"
    fi
}

# Function to cleanup old backups
cleanup_old_backups() {
    echo -e "${YELLOW}Cleaning up old backups...${NC}"
    echo -e "Retention: $RETENTION_COUNT backups or $RETENTION_DAYS days"

    # Remove backups older than retention days
    find "$BACKUP_DIR" -name "mathdb_backup_*.sql.gz" -mtime +$RETENTION_DAYS -delete
    find "$BACKUP_DIR" -name "mathdb_backup_*.sql.gz.sha256" -mtime +$RETENTION_DAYS -delete

    # Keep only the most recent N backups
    ls -t "$BACKUP_DIR"/mathdb_backup_*.sql.gz 2>/dev/null | tail -n +$((RETENTION_COUNT + 1)) | xargs -r rm
    ls -t "$BACKUP_DIR"/mathdb_backup_*.sql.gz.sha256 2>/dev/null | tail -n +$((RETENTION_COUNT + 1)) | xargs -r rm

    # Count remaining backups
    count=$(ls -1 "$BACKUP_DIR"/mathdb_backup_*.sql.gz 2>/dev/null | wc -l)
    echo -e "${GREEN}Cleanup completed. $count backup(s) retained${NC}"
}

# Function to setup automated backups
setup_cron() {
    echo -e "${YELLOW}Setting up automated daily backups...${NC}"

    # Create cron job
    local cron_job="0 2 * * * $PROJECT_ROOT/docker/backup.sh >/dev/null 2>&1"

    # Add to crontab if not exists
    (crontab -l 2>/dev/null | grep -v "$PROJECT_ROOT/docker/backup.sh"; echo "$cron_job") | crontab -

    echo -e "${GREEN}Automated backup scheduled for daily at 2:00 AM${NC}"
    echo -e "${YELLOW}View crontab with: crontab -l${NC}"
}

# Main execution
main() {
    ACTION="backup"

    # Parse arguments
    while [[ $# -gt 0 ]]; do
        case $1 in
            --restore)
                ACTION="restore"
                RESTORE_FILE="$2"
                shift 2
                ;;
            --list)
                ACTION="list"
                shift
                ;;
            --cleanup)
                ACTION="cleanup"
                shift
                ;;
            --setup-cron)
                ACTION="setup-cron"
                shift
                ;;
            *)
                echo -e "${RED}Unknown option: $1${NC}"
                echo "Usage: $0 [--restore <file>] [--list] [--cleanup] [--setup-cron]"
                exit 1
                ;;
        esac
    done

    case $ACTION in
        backup)
            check_container
            create_backup
            cleanup_old_backups
            ;;
        restore)
            check_container
            restore_backup "$RESTORE_FILE"
            ;;
        list)
            list_backups
            ;;
        cleanup)
            cleanup_old_backups
            ;;
        setup-cron)
            setup_cron
            ;;
    esac
}

# Run main function
main "$@"