#!/usr/bin/env python3
"""
Clear stuck tasks from Redis queue after worker restart.

This script identifies tasks that are marked as "active" but have no
corresponding worker processing them, and returns them to the queue.
"""

import redis
import json
import os
import sys
from typing import List, Dict, Any

def get_redis_client():
    """Get Redis client using the same config as the main app."""
    try:
        from modules.redis_config import get_redis_url
        return redis.from_url(get_redis_url())
    except:
        redis_url = os.getenv('REDIS_URL', 'redis://localhost:6379')
        return redis.from_url(redis_url)

def check_stuck_tasks():
    """Check for stuck tasks in Redis."""
    client = get_redis_client()
    
    print("=== REDIS QUEUE STATUS ===")
    processing_count = client.llen("document_processing")
    failed_count = client.llen("document_processing:failed") 
    active_count = client.scard("active_tasks")
    
    print(f"Processing queue: {processing_count} tasks")
    print(f"Failed queue: {failed_count} tasks") 
    print(f"Active tasks: {active_count} tasks")
    
    if active_count > 0:
        print(f"\n⚠️  WARNING: {active_count} tasks marked as 'active' but workers may have been restarted!")
        return True
    elif processing_count > 0:
        print(f"\n✅ {processing_count} tasks waiting in queue (normal)")
        return False
    else:
        print("\n✅ No stuck tasks found")
        return False

def clear_stuck_tasks():
    """Move stuck 'active' tasks back to processing queue."""
    client = get_redis_client()
    
    # Get all active task IDs
    active_tasks = client.smembers("active_tasks")
    if not active_tasks:
        print("No active tasks to clear")
        return
    
    cleared = 0
    for task_id in active_tasks:
        task_id_str = task_id.decode('utf-8')
        try:
            # Try to get the task data
            task_data = client.get(f"task:{task_id_str}")
            if task_data:
                # Parse task to determine priority queue
                task_obj = json.loads(task_data.decode('utf-8'))
                priority = task_obj.get('priority', 'normal')
                
                queue_key = {
                    'high': 'document_processing:high',
                    'normal': 'document_processing', 
                    'low': 'document_processing:low'
                }.get(priority, 'document_processing')
                
                # Move task back to processing queue
                client.lpush(queue_key, task_id_str)
                client.srem("active_tasks", task_id_str)
                cleared += 1
                print(f"✅ Cleared stuck task: {task_id_str} ({task_obj.get('service', 'unknown')}/{task_obj.get('task_type', 'unknown')})")
            else:
                # Task data not found, just remove from active set
                client.srem("active_tasks", task_id_str)
                cleared += 1
                print(f"🗑️  Removed orphaned active task: {task_id_str}")
                
        except Exception as e:
            print(f"❌ Failed to clear task {task_id_str}: {e}")
    
    print(f"\n✅ Cleared {cleared} stuck tasks")

def main():
    """Main function."""
    print("🔍 Checking for stuck tasks in Redis...")
    
    try:
        has_stuck = check_stuck_tasks()
        
        if has_stuck:
            response = input("\n❓ Clear stuck tasks? (y/N): ")
            if response.lower() in ['y', 'yes']:
                print("\n🧹 Clearing stuck tasks...")
                clear_stuck_tasks()
                print("\n✅ Done! Tasks should now be processed by workers.")
            else:
                print("Skipped clearing tasks.")
        
    except Exception as e:
        print(f"❌ Error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()