from celery import shared_task
from django.apps import apps
from django.conf import settings
import requests
import os
import logging
from .ai import extract_metadata_with_llm, validate_date, find_closest_source, find_closest_destination

logger = logging.getLogger(__name__)

@shared_task(bind=True, max_retries=3, default_retry_delay=120, soft_time_limit=800, time_limit=1200)
def send_pdf_to_ocr(self, model_name, obj_id, pdf_field_name):
    """
    Uploads a PDF to the OCR API, processes text through LLM metadata extractor,
    and updates the related Django model.

    Retries ONLY for network-related issues.
    """
    
    app_label = getattr(settings, "OCR_APP_LABEL", "documents")
    Model = apps.get_model(app_label, model_name)

    try:
        obj = Model.objects.get(id=obj_id)
    except Model.DoesNotExist:
        logger.error(f"[OCR TASK] Object {obj_id} of type {model_name} does not exist.")
        return False

    pdf_field = getattr(obj, pdf_field_name)
    pdf_path = pdf_field.path

    # ------------------------------------------------
    # Ensure PDF exists
    # ------------------------------------------------
    if not os.path.exists(pdf_path):
        logger.error(f"[OCR TASK] PDF file missing: {pdf_path}")
        return False

    # ------------------------------------------------
    # 1. Send to OCR API
    # ------------------------------------------------
    try:
        with open(pdf_path, "rb") as f:
            response = requests.post(
                settings.OCR_API_URL,
                files={"file": f},
                data={"rotation": 0, "mode": "Markdown"},
                timeout=1200,
            )
        response.raise_for_status()
    except requests.exceptions.RequestException as exc:
        # Retry only real network failures
        logger.warning(f"[OCR TASK] Network error, retrying: {exc}")
        raise self.retry(exc=exc)

    try:
        ocr_text = response.json().get("text", "")
    except Exception as exc:
        logger.error(f"[OCR TASK] Invalid OCR response JSON: {exc}")
        return False

    # ------------------------------------------------
    # 2. Extract metadata using LLM
    # ------------------------------------------------
    try:
        metadata = extract_metadata_with_llm(model_name, ocr_text)
    except requests.exceptions.RequestException as exc:
        logger.warning(f"[OCR TASK] LLM network failure, retrying: {exc}")
        raise self.retry(exc=exc)
    except Exception as exc:
        logger.error(f"[OCR TASK] Metadata parsing failed: {exc}")
        return False

    # ------------------------------------------------
    # 3. Apply metadata safely
    # ------------------------------------------------
    try:
        # Universal fields
        obj.title = metadata.get("title") or obj.title
        obj.text = metadata.get("body", obj.text)

        # Set number depending on doc type
        if metadata.get("doc_type") == "incoming":
            obj.number = metadata.get("og_document_number") or obj.number
            obj.received_number = metadata.get("document_number") or obj.received_number
        else:
            obj.number = metadata.get("document_number") or obj.number

        # Dates
        doc_date = validate_date(metadata.get("document_date", ""))
        og_date = validate_date(metadata.get("og_document_date", ""))

        if metadata.get("doc_type") == "incoming":
            if og_date:
                obj.og_date = og_date
        else:
            if doc_date:
                obj.date = doc_date

        # Source / Destination logic if needed
        # (This part was commented out in original, keeping it that way or enabling based on user needs. 
        # I'll keep it commented out to match original behavior, but ensure imports are available if they uncomment)
        
        obj.save()

    except Exception as exc:
        logger.error(f"[OCR TASK] Failed to update object {obj_id}: {exc}")
        return False

    logger.info(f"[OCR TASK] Completed successfully for obj {obj_id}")
    return True
