app/worker/tasks.py

import asyncio
import threading
from typing import Any, cast

from app.core.logger import get_logger
from app.models.report import EngineeringReport
from app.services.analysis_service import run_full_analysis
from app.worker.celery_app import celery_app, HIGH_QUEUE, LOW_QUEUE

logger = get_logger(__name__)

# Repos larger than this threshold are routed to the low-priority queue
# so they do not starve small, fast jobs.
LARGE_REPO_MB_THRESHOLD = 50


def route_queue_for_repo(estimated_size_mb: float | None) -> str:
    """Return the appropriate queue name based on estimated repo size."""
    if estimated_size_mb is not None and estimated_size_mb >= LARGE_REPO_MB_THRESHOLD:
        return LOW_QUEUE
    return HIGH_QUEUE


def submit_analysis_task(
    github_url: str,
    base_sha: str | None = None,
    estimated_size_mb: float | None = None,
) -> Any:
    """
    Submit an analysis task to the correct priority queue.

    Call this instead of calling .delay() or .apply_async() directly so
    that queue routing logic stays in one place.
    """
    queue = route_queue_for_repo(estimated_size_mb)
    logger.info(
        "Submitting analysis task",
        extra={"url": github_url, "queue": queue, "size_mb": estimated_size_mb},
    )
    return analyze_repository_task.apply_async(
        args=[github_url, base_sha],
        queue=queue,
    )


def _run_coro_sync(coro: Any) -> Any:
    """
    Run an awaitable to completion from synchronous Celery worker code.

    Uses asyncio.run() when no event loop is running (normal prefork worker).
    Falls back to a thread when a loop is already running (eager test mode).
    """
    try:
        asyncio.get_running_loop()
    except RuntimeError:
        return asyncio.run(coro)

    result: dict[str, Any] = {}

    def _runner() -> None:
        result["value"] = asyncio.run(coro)

    thread = threading.Thread(target=_runner)
    thread.start()
    thread.join()
    return result["value"]


@celery_app.task(name="analyze_repository_task", bind=True, max_retries=2)
def analyze_repository_task(
    self: Any,
    github_url: str,
    base_sha: str | None = None,
) -> dict[str, Any]:
    """
    Celery task wrapper around the async analysis pipeline.

    Runs the full multi-agent pipeline and returns a JSON-serialisable
    EngineeringReport dict. Retries up to 2 times on unexpected failure
    with exponential backoff.
    """
    try:
        report = cast(
            EngineeringReport,
            _run_coro_sync(run_full_analysis(github_url, base_sha)),
        )
        return cast(dict[str, Any], report.model_dump())
    except Exception as exc:
        logger.exception(
            "Analysis task failed",
            extra={"url": github_url, "queue": self.request.delivery_info.get("routing_key"), "error": str(exc)},
        )
        # Exponential backoff: 60s, 120s
        raise self.retry(exc=exc, countdown=60 * (self.request.retries + 1))