import boto3
import redis
import json
import os

def handler(event, context):
    cloudwatch = boto3.client('cloudwatch')
    
    # Connect to Redis
    redis_client = redis.Redis(
        host=os.environ['REDIS_ENDPOINT'],
        port=int(os.environ['REDIS_PORT']),
        decode_responses=True
    )
    
    queue_names = os.environ['CELERY_QUEUES'].split(',')
    total_queue_length = 0
    
    for queue_name in queue_names:
        queue_length = redis_client.llen(queue_name)
        total_queue_length += queue_length
        
        # Publish individual queue metrics
        cloudwatch.put_metric_data(
            Namespace='Celery/Queue',
            MetricData=[
                {
                    'MetricName': f'{queue_name}_length',
                    'Value': queue_length,
                    'Unit': 'Count'
                }
            ]
        )
    
    # Publish total queue length
    cloudwatch.put_metric_data(
        Namespace='Celery/Queue',
        MetricData=[
            {
                'MetricName': 'TotalQueueLength',
                'Value': total_queue_length,
                'Unit': 'Count'
            }
        ]
    )
    
    return {
        'statusCode': 200,
        'body': json.dumps({
            'total_queue_length': total_queue_length,
            'individual_queues': {queue: redis_client.llen(queue) for queue in queue_names}
        })
    }