#!/bin/bash

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
ENVIRONMENT=${1:-local}
BACKUP_DIR="./backups/migrations/$(date +%Y%m%d_%H%M%S)"
ROLLBACK_DIR="./rollback-backups/migrations"

# Determine compose file
if [ "$ENVIRONMENT" == "production" ]; then
    COMPOSE_FILE="docker-compose.production.yml"
    POSTGRES_USER="${POSTGRES_USER:-veza_user}"
    POSTGRES_DB="${POSTGRES_DB:-veza_db}"
else
    COMPOSE_FILE="docker-compose.yml"
    POSTGRES_USER="${POSTGRES_USER:-veza_user}"
    POSTGRES_DB="${POSTGRES_DB:-veza_local}"
fi

# Function to print colored messages
print_info() {
    echo -e "${BLUE}ℹ️  $1${NC}"
}

print_success() {
    echo -e "${GREEN}✅ $1${NC}"
}

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

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

# Function to check prerequisites
check_prerequisites() {
    print_info "Checking prerequisites..."
    
    command -v docker >/dev/null 2>&1 || { 
        print_error "Docker is required but not installed. Aborting."; 
        exit 1; 
    }
    
    command -v docker-compose >/dev/null 2>&1 || { 
        print_error "Docker Compose is required but not installed. Aborting."; 
        exit 1; 
    }
    
    if ! docker info >/dev/null 2>&1; then
        print_error "Docker daemon is not running. Please start Docker and try again.";
        exit 1;
    fi
    
    if [ ! -f "$COMPOSE_FILE" ]; then
        print_error "$COMPOSE_FILE not found. Aborting.";
        exit 1;
    fi
    
    # Check if postgres container is running
    if ! docker-compose -f "$COMPOSE_FILE" ps postgres | grep -q "Up"; then
        print_error "PostgreSQL container is not running. Please start it first.";
        exit 1;
    fi
    
    print_success "Prerequisites check passed"
}

# Function to create backup
create_backup() {
    print_info "Creating database backup..."
    
    mkdir -p "${BACKUP_DIR}"
    mkdir -p "${ROLLBACK_DIR}"
    
    # Backup database
    print_info "Backing up database: ${POSTGRES_DB}..."
    docker-compose -f "$COMPOSE_FILE" exec -T postgres pg_dump -U "${POSTGRES_USER}" "${POSTGRES_DB}" > "${BACKUP_DIR}/database.sql" 2>/dev/null || {
        print_error "Failed to create database backup";
        exit 1;
    }
    
    # Save backup to rollback location
    cp "${BACKUP_DIR}/database.sql" "${ROLLBACK_DIR}/latest.sql"
    
    print_success "Backup created at ${BACKUP_DIR}/database.sql"
}

# Function to check migration status
check_migration_status() {
    print_info "Checking migration status..."
    
    # Check if schema_migrations table exists
    docker-compose -f "$COMPOSE_FILE" exec -T postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -tAc \
        "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = 'public' AND table_name = 'schema_migrations');" \
        > /dev/null 2>&1 || {
        print_warning "schema_migrations table does not exist (will be created)"
        return 0
    }
    
    # Get list of applied migrations
    APPLIED_MIGRATIONS=$(docker-compose -f "$COMPOSE_FILE" exec -T postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -tAc \
        "SELECT COUNT(*) FROM schema_migrations;" 2>/dev/null || echo "0")
    
    print_info "Currently applied migrations: ${APPLIED_MIGRATIONS}"
}

# Function to run backend API migrations
run_backend_migrations() {
    print_info "Running Backend API migrations (GORM AutoMigrate)..."
    
    # Check if backend-api container is running
    if ! docker-compose -f "$COMPOSE_FILE" ps backend-api | grep -q "Up"; then
        print_warning "Backend API container is not running. Starting it..."
        docker-compose -f "$COMPOSE_FILE" up -d backend-api
        sleep 10
    fi
    
    # Trigger migrations by restarting the backend-api service
    # The backend API runs migrations automatically on startup via Database.Initialize()
    print_info "Restarting backend-api to trigger migrations..."
    docker-compose -f "$COMPOSE_FILE" restart backend-api || {
        print_error "Failed to restart backend-api";
        return 1;
    }
    
    # Wait for backend to be ready
    print_info "Waiting for backend-api to be ready..."
    sleep 15
    
    # Check if backend is healthy
    if curl -sf "http://localhost:8080/health" >/dev/null 2>&1; then
        print_success "Backend API migrations completed successfully"
        return 0
    else
        print_warning "Backend API health check failed (migrations may still be running)"
        return 0
    fi
}

# Function to run chat server migrations
run_chat_server_migrations() {
    print_info "Checking for Chat Server migrations..."
    
    # Check if chat-server has migrations directory
    if [ ! -d "veza-chat-server/migrations" ]; then
        print_warning "Chat Server migrations directory not found, skipping"
        return 0
    fi
    
    # Check if chat-server uses sqlx migrate
    if docker-compose -f "$COMPOSE_FILE" exec -T chat-server cargo --version >/dev/null 2>&1; then
        print_info "Running Chat Server migrations (sqlx migrate)..."
        
        # Run sqlx migrate (if available in the container)
        docker-compose -f "$COMPOSE_FILE" exec -T chat-server \
            sqlx migrate run --database-url "${DATABASE_URL}" 2>/dev/null || {
            print_warning "sqlx migrate not available or migrations already applied"
            return 0
        }
        
        print_success "Chat Server migrations completed"
    else
        print_warning "Chat Server migrations will be handled by the service on startup"
    fi
}

# Function to verify migrations
verify_migrations() {
    print_info "Verifying migrations..."
    
    # Check database connection
    if ! docker-compose -f "$COMPOSE_FILE" exec -T postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -c "SELECT 1;" >/dev/null 2>&1; then
        print_error "Failed to connect to database";
        return 1;
    fi
    
    # Check if schema_migrations table exists
    TABLE_EXISTS=$(docker-compose -f "$COMPOSE_FILE" exec -T postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -tAc \
        "SELECT EXISTS (SELECT FROM information_schema.tables WHERE table_schema = 'public' AND table_name = 'schema_migrations');" 2>/dev/null || echo "false")
    
    if [ "$TABLE_EXISTS" = "t" ]; then
        MIGRATION_COUNT=$(docker-compose -f "$COMPOSE_FILE" exec -T postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" -tAc \
            "SELECT COUNT(*) FROM schema_migrations;" 2>/dev/null || echo "0")
        print_info "Migrations table exists with ${MIGRATION_COUNT} applied migrations"
    fi
    
    print_success "Migration verification passed"
    return 0
}

# Function to rollback
rollback() {
    print_error "Migration failed! Initiating rollback..."
    
    if [ ! -f "${ROLLBACK_DIR}/latest.sql" ]; then
        print_error "No rollback backup found. Manual intervention required."
        exit 1
    fi
    
    print_info "Rolling back database to previous state..."
    
    # Drop and recreate database (or restore from backup)
    print_warning "Restoring database from backup..."
    
    # Restore database
    docker-compose -f "$COMPOSE_FILE" exec -T postgres psql -U "${POSTGRES_USER}" -d "${POSTGRES_DB}" < "${ROLLBACK_DIR}/latest.sql" || {
        print_error "Failed to restore database from backup. Manual intervention required.";
        exit 1;
    }
    
    print_warning "Rollback completed. Please verify database manually."
    exit 1
}

# Main migration flow
main() {
    print_info "Starting database migrations for ${ENVIRONMENT} environment..."
    
    # Trap errors and rollback
    trap rollback ERR
    
    check_prerequisites
    check_migration_status
    create_backup
    run_backend_migrations
    run_chat_server_migrations
    
    # Wait a bit for all migrations to complete
    sleep 5
    
    # Verify migrations
    if ! verify_migrations; then
        rollback
    fi
    
    # Disable trap on success
    trap - ERR
    
    print_success "Database migrations completed successfully for ${ENVIRONMENT}!"
    print_info "Backup location: ${BACKUP_DIR}/database.sql"
}

# Run main function
main