from celery import Celery
from kombu import Queue


def init_celery(celery, app):
    celery_config = {
        'task_serializer': app.config.get('CELERY_TASK_SERIALIZER', 'json'),
        'result_serializer': app.config.get('CELERY_RESULT_SERIALIZER', 'json'),
        'accept_content': app.config.get('CELERY_ACCEPT_CONTENT', ['json']),
        'timezone': app.config.get('CELERY_TIMEZONE', 'UTC'),
        'enable_utc': app.config.get('CELERY_ENABLE_UTC', True),
    }
    celery.conf.update(**celery_config)

    # Setting up Celery task queues
    celery.conf.task_queues = (
        Queue('embeddings', routing_key='embeddings.key', queue_arguments={'x-max-priority': 10}),
        Queue('llm_interactions', routing_key='llm_interactions.key', queue_arguments={'x-max-priority': 5}),
    )

    # Ensuring tasks execute with Flask application context
    class ContextTask(celery.Task):
        def __call__(self, *args, **kwargs):
            with app.app_context():
                return self.run(*args, **kwargs)

    celery.Task = ContextTask


def make_celery(app_name, config):
    celery = Celery(app_name, broker=config['CELERY_BROKER_URL'], backend=config['CELERY_RESULT_BACKEND'])
    return celery