certificate_upload_module.py

"""Upload generated certificates to the public Build Small gallery dataset."""
import os
import time
import tempfile
import logging

from PIL import Image as PILImage
from datasets import Dataset, Image, load_dataset, concatenate_datasets
from huggingface_hub import HfApi

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# PUBLIC dataset that acts as the gallery of issued certificates
CERTIFICATE_DATASET_NAME = "build-small-hackathon/build-small-certificates"

HF_TOKEN = os.getenv("HF_TOKEN")


def safe_add_certificate_to_dataset(certificate_image, hf_username, max_retries=5, retry_delay=3):
    """Append a certificate image to the dataset, handling empty/existing datasets and dedup."""
    try:
        if not hf_username or not hf_username.strip():
            return False, "❌ Error: HF username is required"
        if certificate_image is None:
            return False, "❌ Error: Certificate image is required"
        hf_username = hf_username.strip()
        logger.info(f"Processing certificate for user: {hf_username}")

        existing_dataset = None
        load_successful = False
        is_empty_dataset = False

        for attempt in range(max_retries):
            try:
                existing_dataset = load_dataset(CERTIFICATE_DATASET_NAME, split="train", token=HF_TOKEN)
                logger.info(f"Loaded {len(existing_dataset)} existing certificates")
                load_successful = True
                break
            except Exception as load_error:
                error_str = str(load_error).lower()
                if "corresponds to no data" in error_str or "no data" in error_str or "doesn't exist" in error_str or "not found" in error_str:
                    logger.info("Dataset empty / not yet created — will create first entry")
                    is_empty_dataset = True
                    load_successful = True
                    existing_dataset = None
                    break
                logger.warning(f"Attempt {attempt + 1} failed: {str(load_error)[:120]}")
                if attempt < max_retries - 1:
                    time.sleep(retry_delay)

        if not load_successful:
            return False, ("❌ Certificate upload temporarily unavailable. Please try again in a few minutes.")

        # Dedup by username (stored in the 'label' column)
        if existing_dataset is not None:
            if hf_username in existing_dataset["label"]:
                return True, f"a certificate for '{hf_username}' already exists in the gallery."

        with tempfile.TemporaryDirectory() as temp_dir:
            if isinstance(certificate_image, PILImage.Image):
                temp_image_path = os.path.join(temp_dir, f"certificate_{hf_username}_{int(time.time())}.png")
                certificate_image.save(temp_image_path, "PNG")
            elif isinstance(certificate_image, str) and os.path.exists(certificate_image):
                temp_image_path = certificate_image
            else:
                return False, "❌ Error: Invalid image format provided"

            new_dataset = Dataset.from_dict(
                {"image": [temp_image_path], "label": [hf_username]}
            ).cast_column("image", Image())

            if existing_dataset is not None and not is_empty_dataset:
                combined_dataset = concatenate_datasets([existing_dataset, new_dataset])
            else:
                combined_dataset = new_dataset

            try:
                combined_dataset.push_to_hub(CERTIFICATE_DATASET_NAME, private=False, token=HF_TOKEN)
                logger.info(f"Saved certificate. Total now: {len(combined_dataset)}")
                return True, f"✅ saved to the gallery for {hf_username}."
            except Exception as upload_error:
                msg = str(upload_error).lower()
                if any(i in msg for i in ("rate limit", "429", "too many requests")):
                    return False, "⏳ Upload busy due to high load — please try again in 10–15 minutes."
                logger.error(f"Upload failed: {upload_error}")
                return False, f"❌ Certificate upload failed: {str(upload_error)}"
    except Exception as e:
        logger.error(f"Unexpected error in certificate upload: {e}")
        return False, f"❌ Certificate upload failed: {str(e)}"


def upload_user_certificate(certificate_image, hf_username):
    """Public entry point — returns (success, message)."""
    if not certificate_image:
        return False, "❌ No certificate image provided"
    if not hf_username or not hf_username.strip():
        return False, "❌ HF username is required"
    return safe_add_certificate_to_dataset(certificate_image, hf_username)