nodes/cropper.py

"""execute_crops node — Gemini code_execution for agentic cropping (PoC 1 style)."""
from __future__ import annotations

import io
import logging
import time
import uuid
from collections.abc import Callable
from concurrent.futures import ThreadPoolExecutor, as_completed

from google import genai
from google.genai import types
from PIL import Image

from config import CROPPER_MODEL, GOOGLE_API_KEY
from prompts.cropper import CROPPER_PROMPT_TEMPLATE
from state import CropTask, DrawingReaderState, ImageRef
from tools.crop_cache import CropCache
from tools.image_store import ImageStore
from tools.pdf_processor import get_page_image_bytes

logger = logging.getLogger(__name__)

# Type alias for the progress callback.
# Signature: (completed_ref, crop_task, source, completed_count, total_count)
ProgressCallback = Callable[[ImageRef, CropTask, str, int, int], None]

# Retry settings for transient API errors (429 / 503)
MAX_RETRIES = 3
RETRY_BASE_DELAY = 2.0  # seconds


def _extract_last_image(response) -> Image.Image | None:
    """Extract the last generated image from a Gemini code_execution response."""
    last_image = None
    for part in response.candidates[0].content.parts:
        # Try as_image() first
        try:
            img_data = part.as_image()
            if img_data is not None:
                last_image = Image.open(io.BytesIO(img_data.image_bytes))
                continue
        except Exception:
            pass
        # Fallback: inline_data
        try:
            if hasattr(part, "inline_data") and part.inline_data is not None:
                img_bytes = part.inline_data.data
                last_image = Image.open(io.BytesIO(img_bytes))
        except Exception:
            pass
    return last_image


def _execute_single_crop_sync(
    client: genai.Client,
    page_image_bytes: bytes,
    crop_task: CropTask,
    image_store: ImageStore,
) -> tuple[ImageRef, bool]:
    """Execute one crop via Gemini code_execution (synchronous).

    Includes retry logic for transient 503/429 errors.

    Returns
    -------
    (image_ref, is_fallback)
        ``is_fallback`` is True when Gemini failed to produce a crop and the
        full page image was returned instead.  Fallbacks should NOT be cached.
    """
    prompt = CROPPER_PROMPT_TEMPLATE.format(
        crop_instruction=crop_task["crop_instruction"],
    )

    image_part = types.Part.from_bytes(data=page_image_bytes, mime_type="image/png")

    # Retry loop for transient API errors
    response = None
    for attempt in range(MAX_RETRIES):
        try:
            response = client.models.generate_content(
                model=CROPPER_MODEL,
                contents=[image_part, prompt],
                config=types.GenerateContentConfig(
                    tools=[types.Tool(code_execution=types.ToolCodeExecution)]
                ),
            )
            break
        except Exception as e:
            err_str = str(e)
            if ("503" in err_str or "429" in err_str or "UNAVAILABLE" in err_str):
                delay = RETRY_BASE_DELAY * (2 ** attempt)
                logger.warning(
                    "Crop API error (attempt %d/%d): %s — retrying in %.1fs",
                    attempt + 1, MAX_RETRIES, err_str[:120], delay,
                )
                time.sleep(delay)
            else:
                raise

    is_fallback = True
    if response is not None:
        final_image = _extract_last_image(response)
        if final_image is not None:
            is_fallback = False
        else:
            final_image = Image.open(io.BytesIO(page_image_bytes))
    else:
        # All retries exhausted
        final_image = Image.open(io.BytesIO(page_image_bytes))

    crop_id = f"crop_{uuid.uuid4().hex[:6]}"
    ref = image_store.save_crop(
        page_num=crop_task["page_num"],
        crop_id=crop_id,
        image=final_image,
        label=crop_task["label"],
    )
    return ref, is_fallback


def execute_crops(
    state: DrawingReaderState,
    image_store: ImageStore,
    crop_cache: CropCache | None = None,
    progress_callback: ProgressCallback | None = None,
) -> dict:
    """Execute all crop tasks concurrently, reusing cached crops when possible.

    Parameters
    ----------
    progress_callback
        Optional callback invoked on the **main thread** each time a crop
        completes (or is served from cache).  Called with
        ``(image_ref, crop_task, source, completed_count, total_count)``
        where *source* is ``"cached"``, ``"completed"``, or ``"fallback"``.
    """
    crop_tasks = state.get("crop_tasks", [])
    page_image_dir = state["page_image_dir"]

    if not crop_tasks:
        return {"status_message": ["No crop tasks to execute."]}

    total_count = len(crop_tasks)
    completed_count = 0

    # ----- Phase 1: Separate cache hits from tasks that need API calls -----
    image_refs: list[ImageRef] = []         # final ordered results
    tasks_to_execute: list[tuple[int, CropTask]] = []  # (original_index, task)
    cache_hits = 0

    for i, ct in enumerate(crop_tasks):
        if crop_cache is not None:
            cached_ref = crop_cache.lookup(ct["page_num"], ct["crop_instruction"])
            if cached_ref is not None:
                image_refs.append(cached_ref)
                cache_hits += 1
                completed_count += 1
                logger.info(
                    "Reusing cached crop for '%s' (page %d)",
                    ct["label"], ct["page_num"],
                )
                # Notify the UI immediately for each cache hit
                if progress_callback is not None:
                    progress_callback(
                        cached_ref, ct, "cached", completed_count, total_count,
                    )
                continue
        # Not cached — needs an API call
        tasks_to_execute.append((i, ct))

    # ----- Phase 2: Execute uncached crops via Gemini -----
    errors: list[str] = []

    if tasks_to_execute:
        client = genai.Client(api_key=GOOGLE_API_KEY)

        with ThreadPoolExecutor(max_workers=min(len(tasks_to_execute), 4)) as pool:
            future_to_idx: dict = {}
            for exec_idx, (_, ct) in enumerate(tasks_to_execute):
                page_bytes = get_page_image_bytes(page_image_dir, ct["page_num"])
                future = pool.submit(
                    _execute_single_crop_sync, client, page_bytes, ct, image_store,
                )
                future_to_idx[future] = exec_idx

            # Process results as they arrive — this runs on the MAIN thread,
            # so we can safely invoke the Streamlit progress callback here.
            for future in as_completed(future_to_idx):
                exec_idx = future_to_idx[future]
                orig_idx, ct = tasks_to_execute[exec_idx]
                try:
                    ref, is_fallback = future.result()
                    image_refs.append(ref)
                    completed_count += 1

                    # Register in cache (only successful targeted crops)
                    if crop_cache is not None:
                        crop_cache.register(
                            page_num=ct["page_num"],
                            crop_instruction=ct["crop_instruction"],
                            label=ct["label"],
                            image_ref=ref,
                            is_fallback=is_fallback,
                        )

                    # Notify the UI as each crop completes
                    if progress_callback is not None:
                        source = "fallback" if is_fallback else "completed"
                        progress_callback(
                            ref, ct, source, completed_count, total_count,
                        )

                except Exception as e:
                    completed_count += 1
                    errors.append(f"Crop task {orig_idx} failed: {e}")
                    logger.error("Crop task %d failed: %s", orig_idx, e)

    # ----- Phase 3: Build status message -----
    api_count = len(tasks_to_execute) - len(errors)
    parts = [f"Completed {len(image_refs)} of {total_count} crops"]
    if cache_hits:
        parts.append(f"({cache_hits} from cache, {api_count} new)")
    if errors:
        parts.append(f"Errors: {'; '.join(errors)}")
    status = ". ".join(parts) + "."

    if crop_cache is not None:
        logger.info(crop_cache.stats)

    return {
        "image_refs": image_refs,
        "status_message": [status],
    }