"""Celery tasks.

`process_document_task` is the single async entrypoint: the API enqueues
one task per upload, the worker pulls the blob, runs the orchestrator, and
writes the result back to the `jobs` table.

Tasks must be idempotent at the boundary: if the worker crashes mid-OCR,
the next retry should pick up the same blob and produce the same result.
We therefore re-fetch the row by id on every transition rather than
threading state through closures.
"""

from __future__ import annotations

from uuid import UUID

from celery.exceptions import Reject

from ocr_sprint.db.base import session_scope
from ocr_sprint.db.repositories import JobRepository
from ocr_sprint.pipeline.orchestrator import run_pipeline
from ocr_sprint.storage.blob import get_blob_storage
from ocr_sprint.utils.logging import get_logger
from ocr_sprint.worker.celery_app import celery_app

_logger = get_logger(__name__)


@celery_app.task(name="ocr_sprint.process_document", bind=True, max_retries=0)  # type: ignore[untyped-decorator]
def process_document_task(self: object, job_id_str: str) -> str:
    """Run the OCR pipeline for `job_id` and persist the structured result.

    Returns the final job status as a string so callers can wait on
    `AsyncResult.get()` if they want to (mainly tests).
    """
    job_id = UUID(job_id_str)
    log = _logger.bind(job_id=job_id_str)
    storage = get_blob_storage()

    with session_scope() as session:
        repo = JobRepository(session)
        row = repo.get(job_id)
        if row is None:
            raise Reject(f"Job {job_id_str} not found", requeue=False)
        repo.mark_processing(job_id)
        blob_key = row.blob_key

    if not blob_key:
        with session_scope() as session:
            JobRepository(session).mark_failed(job_id, error="missing blob_key")
        return "failed"

    try:
        content = storage.get(blob_key)
    except FileNotFoundError as exc:
        log.error("worker.blob_missing", error=str(exc))
        with session_scope() as session:
            JobRepository(session).mark_failed(job_id, error=f"blob missing: {exc}")
        return "failed"

    try:
        output = run_pipeline(content)
    except Exception as exc:
        log.exception("worker.pipeline_error")
        with session_scope() as session:
            JobRepository(session).mark_failed(job_id, error=str(exc))
        return "failed"

    flags = [f.value for f in output.result.review_flags]
    log.info(
        "worker.completed",
        status=output.status.value,
        confidence=round(output.confidence, 3),
        flags=flags,
    )
    with session_scope() as session:
        JobRepository(session).mark_completed(
            job_id,
            status=output.status,
            confidence=output.confidence,
            result=output.result.model_dump(mode="json"),
            review_flags=flags,
        )
    return output.status.value
