import os
from celery import Celery

# Use Redis as the broker and result backend
REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379/0")

celery_app = Celery(
    "chatbot_tasks",
    broker=REDIS_URL,
    backend=REDIS_URL
)

celery_app.conf.update(
    task_serializer="json",
    accept_content=["json"],
    result_serializer="json",
    timezone="UTC",
    enable_utc=True,
    # Crucial for preventing local AI overload: 
    # worker_concurrency limits how many simultaneous Ollama invocations can occur.
    worker_concurrency=2,
)

@celery_app.task(bind=True, max_retries=3)
def process_ai_request_task(self, query: str, context: str, custom_prompt: str = None):
    """
    Background task wrapper for the enterprise AI fallback chain.
    If Local AI freezes, Celery can enforce hard time limits here.
    """
    import asyncio
    import services.ai_service_enterprise as enterprise_ai
    
    loop = asyncio.get_event_loop()
    try:
        # Enforce a strict timeout on local generation
        result = loop.run_until_complete(
            asyncio.wait_for(
                enterprise_ai.generate_chatbot_response(query, context, custom_prompt),
                timeout=30.0
            )
        )
        return result
    except asyncio.TimeoutError:
        # If it times out, the fallback chain should technically handle it, 
        # but if the node freezes entirely, this rescues the worker.
        return {
            "response": "Request timed out during AI generation.",
            "provider": "error",
            "tokens_used": 0
        }
    except Exception as exc:
        raise self.retry(exc=exc, countdown=5)