src/tasks/ocr_tasks.py

"""
TradeFlow AI — Wired OCR & Processing Tasks (Phase 2+3 implementation)
"""

from __future__ import annotations

import structlog

from .celery_app import celery_app

log = structlog.get_logger()


async def _load_batch_context(batch_id: str) -> dict:
    """Load batch and document rows for the LangGraph worker."""
    from supabase import acreate_client

    from ..config import settings

    supabase = await acreate_client(settings.SUPABASE_URL, settings.SUPABASE_SERVICE_KEY.get_secret_value())
    try:
        batch_res = await (
            supabase.table("batches").select("*").eq("id", batch_id).single().execute()
        )
        docs_res = await supabase.table("documents").select("*").eq("batch_id", batch_id).execute()
        batch = batch_res.data or {}
        documents = [
            {
                "doc_id": row["id"],
                "doc_type": row.get("doc_type"),
                "storage_path": row.get("storage_path"),
                "original_name": row.get("original_name"),
                "pages": [],
                "extracted_data": None,
                "quality_score": float(row.get("quality_score") or 0.0),
                "ocr_method": row.get("ocr_engine_used"),
                "error": row.get("error_message"),
                "ocr_candidates": {},
                "ocr_conflicts": [],
                "field_confidences": {},
            }
            for row in docs_res.data
        ]
        return {"batch": batch, "documents": documents}
    finally:
        pass


async def _persist_graph_result(batch_id: str, result: dict) -> None:
    """Persist OCR/extraction confidence so dashboard/eval can measure accuracy."""
    from supabase import acreate_client

    from ..config import settings

    supabase = await acreate_client(settings.SUPABASE_URL, settings.SUPABASE_SERVICE_KEY.get_secret_value())
    try:
        await supabase.table("extracted_fields").delete().eq("batch_id", batch_id).execute()
        await supabase.table("validation_results").delete().eq("batch_id", batch_id).execute()

        for doc in result.get("documents", []):
            await supabase.table("documents").update({
                "quality_score": doc.get("quality_score"),
                "ocr_engine_used": doc.get("ocr_method"),
                "overall_ocr_confidence": _average_confidence(doc.get("field_confidences") or {}),
                "error_message": doc.get("error"),
                "status": "ocr_complete" if not doc.get("error") else "error",
            }).eq("id", doc["doc_id"]).execute()

            extracted_data = doc.get("extracted_data") or {}
            field_confidences = doc.get("field_confidences") or {}
            rows = [
                {
                    "batch_id": batch_id,
                    "document_id": doc["doc_id"],
                    "ceisa_field": field,
                    "raw_ocr_value": str(value),
                    "extracted_value": str(value),
                    "normalized_value": str(value),
                    "confidence": float(field_confidences.get(field, 0.0)),
                    "extraction_method": "direct_ocr",
                }
                for field, value in extracted_data.items()
            ]
            if rows:
                await supabase.table("extracted_fields").insert(rows).execute()

        for validation in result.get("validation_results", []):
            await supabase.table("validation_results").insert({
                "batch_id": batch_id,
                "rule_id": validation.get("rule_id", "UNKNOWN"),
                "rule_name": validation.get("rule_name", validation.get("message", "Validation")),
                "severity": validation.get("severity", "WARNING"),
                "error_message": validation.get("message"),
                "affected_fields": validation.get("affected_fields", []),
            }).execute()

        status = "review_ready" if result.get("needs_human_review") else "validated"
        await supabase.table("batches").update({
            "status": status,
            "risk_level": result.get("risk_level"),
            "customs_readiness_score": result.get("customs_readiness_score", result.get("_crs_score")),
            "crs_grade": result.get("crs_grade", result.get("_crs_grade")),
            "rejection_probability": result.get("rejection_probability", result.get("_rejection_prob")),
            "langgraph_thread_id": batch_id,
        }).eq("id", batch_id).execute()
    finally:
        pass


async def _update_batch_status(batch_id: str, status: str, error_message: str | None = None) -> None:
    """Best-effort status update for the upload/detail UI."""
    from supabase import acreate_client

    from ..config import settings

    supabase = await acreate_client(settings.SUPABASE_URL, settings.SUPABASE_SERVICE_KEY.get_secret_value())
    payload = {"status": status}
    try:
        await supabase.table("batches").update(payload).eq("id", batch_id).execute()
    except Exception as exc:
        log.warning(
            "Failed to update batch status",
            batch_id=batch_id,
            status=status,
            pipeline_error=error_message,
            error=str(exc),
        )


def _average_confidence(confidences: dict) -> float:
    values = [float(value) for value in confidences.values()]
    return round(sum(values) / len(values), 4) if values else 0.0


def run_preprocess_pipeline_sync(batch_id: str, mark_error_on_failure: bool = False) -> None:
    """Run the extraction pipeline in-process when Celery is unavailable."""
    from ..ai.graph import extraction_graph
    from .celery_app import get_worker_loop

    loop = get_worker_loop()
    try:
        config = {"configurable": {"thread_id": batch_id}}
        loop.run_until_complete(_update_batch_status(batch_id, "ocr_running"))
        context = loop.run_until_complete(_load_batch_context(batch_id))
        initial_state = {
            "batch_id": batch_id,
            "company_id": context["batch"].get("company_id") or "",
            "documents": context["documents"],
            "combined_data": {},
            "validation_results": [],
            "needs_human_review": False,
            "risk_level": "UNKNOWN",
            "customs_readiness_score": None,
            "crs_grade": None,
            "rejection_probability": None,
            "risk_features": {},
            "ocr_conflicts": [],
            "field_confidences": {},
            "steps": [],
        }
        result = loop.run_until_complete(
            extraction_graph.ainvoke(initial_state, config=config)
        )
        loop.run_until_complete(_persist_graph_result(batch_id, result))
    except Exception as exc:
        if mark_error_on_failure:
            loop.run_until_complete(_update_batch_status(batch_id, "error", str(exc)))
        raise


@celery_app.task(bind=True, queue="high", max_retries=3, default_retry_delay=10)
def preprocess_document(self, batch_id: str) -> None:
    """Entry point: kicks off LangGraph extraction pipeline."""
    log.info("Starting extraction pipeline", batch_id=batch_id)
    try:
        run_preprocess_pipeline_sync(batch_id)
        log.info("Extraction pipeline complete", batch_id=batch_id)
    except Exception as exc:
        log.error("Pipeline failed", batch_id=batch_id, error=str(exc))
        if self.request.retries >= self.max_retries:
            from .celery_app import get_worker_loop

            loop = get_worker_loop()
            loop.run_until_complete(_update_batch_status(batch_id, "error", str(exc)))
        raise self.retry(exc=exc)


@celery_app.task(bind=True, queue="high")
def run_ocr(self, batch_id: str) -> None:
    log.info("run_ocr delegated to LangGraph", batch_id=batch_id)


@celery_app.task(bind=True, queue="high")
def extract_fields(self, batch_id: str) -> None:
    log.info("extract_fields delegated to LangGraph", batch_id=batch_id)


@celery_app.task(bind=True, queue="default")
def validate_fields(self, batch_id: str) -> None:
    log.info("validate_fields delegated to LangGraph", batch_id=batch_id)


@celery_app.task(bind=True, queue="default")
def recommend_hs(self, batch_id: str, product_description: str) -> list[dict]:
    """Run HS code recommendation for a batch."""
    import asyncio

    from ..services.hs_svc import hs_recommend_service
    log.info("Running HS recommendation", batch_id=batch_id)
    return asyncio.get_event_loop().run_until_complete(
        hs_recommend_service.recommend(product_description)
    )


@celery_app.task(bind=True, queue="default")
def assess_risk(self, batch_id: str) -> dict:
    """Compute CRS and persist to DB."""
    from ..services.predictor_svc import rejection_predictor
    log.info("Assessing risk", batch_id=batch_id)
    # Stub features — in full impl these come from graph state persisted in Redis
    features = {
        "doc_quality_score": 0.95,
        "completeness_score": 0.88,
        "consistency_score": 1.0,
        "historical_rate": 0.80,
        "hs_confidence": 0.85,
        "cif_value_usd": 15000.0,
        "package_count": 10,
        "gross_weight_kg": 500.0,
    }
    rejection_predictor.load()
    rejection_prob = rejection_predictor.predict_proba(features)
    return {"rejection_prob": rejection_prob, "batch_id": batch_id}