from django.tasks import task
from submissions.models import ImageCrop, SubmissionFile, OCRStatus

from PIL import Image
import pytesseract


def _process_ocr_file(obj, file, crop: ImageCrop | None = None) -> bool:
    if obj.ocr_status != OCRStatus.TO_PROCESS:
        return False

    obj.ocr_status = OCRStatus.PROCESSING
    obj.save(update_fields=["ocr_status"])

    try:
        image = Image.open(file)

        if crop:
            image = image.crop((crop.x1, crop.y1, crop.x2, crop.y2))

        obj.ocr_raw_data = pytesseract.image_to_string(image)
        obj.ocr_status = OCRStatus.PROCESSED

    except Exception as e:
        obj.ocr_status = OCRStatus.ERROR
        raise e

    finally:
        obj.save(update_fields=["ocr_status", "ocr_raw_data"])

    return True


@task
def process_submission_files():
    """Process uploaded submission files (OCR, validation, etc.)"""
    for file in SubmissionFile.objects.filter(ocr_status=OCRStatus.TO_PROCESS):
        process_submission_file.enqueue(str(file.id))


@task
def process_submission_file(file_id: str):
    file = SubmissionFile.objects.filter(id=file_id).first()

    if not file:
        return

    _process_ocr_file(file, file.file)

    process_crops.enqueue(file_id)


@task
def process_crops(file_id: str):
    file = SubmissionFile.objects.filter(id=file_id).first()

    if not file:
        return

    for crop in file.crops.all():
        process_crop.enqueue(str(crop.id))


@task
def process_crop(crop_id: str):
    crop = ImageCrop.objects.filter(id=crop_id).first()

    if not crop:
        return

    _process_ocr_file(crop, crop.submission_file.file, crop=crop)