ai_selector.py

"""
AI-based photo selection using OpenAI Vision.
Sends images in batches, asks for the best ones by quality criteria, returns filenames as JSON.
"""

import base64
import json
import re
import time
from pathlib import Path

from openai import OpenAI
from config import OPENAI_MAX_RETRIES, OPENAI_MODEL


# -----------------------------------------------------------------------------
# Prompt sent to OpenAI Vision (instructs model to return ONLY JSON)
# -----------------------------------------------------------------------------
SELECTION_PROMPT = """You are a professional photo curator. Look at the images provided.

For each image you receive, consider:
- Sharpness: Is the image sharp and clear (not blurry)?
- Focus: Is the subject in focus?
- Lighting: Is the lighting flattering and well-balanced?
- Composition: Is the framing and composition strong (rule of thirds, balance, no cut-off subjects)?
- No duplicates: If two images are nearly identical, pick only the better one.
- Best overall aesthetic: Which images would you keep in a final album?

Your task: From the images in this batch, select ONLY the best photos. Return the filenames of selected photos only.

You MUST respond with valid JSON only, no other text. Use this exact format:
{
  "selected": ["filename1.jpg", "filename2.jpg"]
}

If a photo's filename is not visible, use the order: the first image is "image_1.jpg", second is "image_2.jpg", etc., and list those names in "selected".

Return ONLY the JSON object. No markdown, no explanation, no code block."""

# Prompt when we want only the single best photo from the batch
SELECTION_PROMPT_ONE = """You are a professional photo curator. Look at the images provided.

For each image, consider: Sharpness, Focus, Lighting, Composition, and overall aesthetic.

Your task: From this batch, select exactly ONE photo—the single best. Return only that filename.

You MUST respond with valid JSON only, no other text. Use this exact format:
{
  "selected": ["filename.jpg"]
}

If a photo's filename is not visible, use the order: the first image is "image_1.jpg", second is "image_2.jpg", etc. Put exactly one name in "selected".

Return ONLY the JSON object. No markdown, no explanation, no code block."""


def _encode_image(path: Path) -> str:
    """Read image file and return base64 string for API."""
    with open(path, "rb") as f:
        return base64.standard_b64encode(f.read()).decode("utf-8")


def _guess_mime(path: Path) -> str:
    """Return MIME type based on extension for OpenAI API."""
    ext = path.suffix.lower()
    mime = {
        ".jpg": "image/jpeg",
        ".jpeg": "image/jpeg",
        ".png": "image/png",
        ".gif": "image/gif",
        ".webp": "image/webp",
        ".heic": "image/heic",
        ".bmp": "image/bmp",
        ".tiff": "image/tiff",
        ".tif": "image/tiff",
    }
    return mime.get(ext, "image/jpeg")


def _parse_response_json(text: str, image_paths: list[Path]) -> list[str]:
    """
    Parse model response into list of selected filenames.
    Handles optional markdown code block and maps image_1.jpg style back to real filenames if used.
    """
    # Strip markdown code block if present
    stripped = text.strip()
    if "```" in stripped:
        lines = stripped.split("\n")
        lines = [l for l in lines if not l.strip().startswith("```")]
        stripped = "\n".join(lines)
    stripped = stripped.strip()

    # Try parsing full response first
    data = None
    try:
        data = json.loads(stripped)
    except json.JSONDecodeError:
        # Try to find JSON object in response (e.g. {"selected": ["a.jpg", "b.jpg"]})
        match = re.search(r'\{[^{}]*"selected"\s*:\s*\[[^\]]*\]\s*\}', stripped, re.DOTALL)
        if match:
            try:
                data = json.loads(match.group(0))
            except json.JSONDecodeError:
                pass

    if not data or "selected" not in data:
        return []
    selected = data.get("selected") or []

    # Map placeholder names (image_1.jpg, image_2.jpg, ...) to actual filenames
    result = []
    for name in selected:
        if not isinstance(name, str):
            name = str(name).strip()
        else:
            name = name.strip()
        lower = name.lower()
        if lower.startswith("image_") and (".jpg" in lower or ".jpeg" in lower or ".png" in lower):
            try:
                num = int("".join(c for c in name if c.isdigit()) or "0")
                if 1 <= num <= len(image_paths):
                    name = image_paths[num - 1].name
            except (ValueError, IndexError):
                pass
        result.append(name)
    return result


def select_best_photos(
    image_paths: list[str] | list[Path],
    max_selections: int | None = None,
) -> list[str]:
    """
    Send the given images to OpenAI Vision and return the list of selected filenames only.

    - image_paths: list of paths (str or Path) to image files.
    - max_selections: if set (e.g. 1), only return up to this many filenames; AI prompt asks for one best.
    - Returns: list of filenames (e.g. ["IMG_1234.jpg"]) that the AI selected.
    """
    paths = [Path(p) for p in image_paths]
    if not paths:
        return []

    prompt = SELECTION_PROMPT_ONE if max_selections == 1 else SELECTION_PROMPT
    client = OpenAI()

    # Build content: prompt + one image per message part (with filename in detail for reference)
    content = [{"type": "text", "text": prompt}]
    for i, p in enumerate(paths):
        if not p.is_file():
            continue
        content.append({
            "type": "image_url",
            "image_url": {
                "url": f"data:{_guess_mime(p)};base64,{_encode_image(p)}",
                "detail": "low",
            },
        })
        # Add a short text hint so the model can map back to filename
        content.append({
            "type": "text",
            "text": f"[Image {i + 1} filename: {p.name}]",
        })

    last_error = None
    for attempt in range(OPENAI_MAX_RETRIES):
        try:
            response = client.chat.completions.create(
                model=OPENAI_MODEL,
                messages=[{"role": "user", "content": content}],
                max_tokens=1024,
            )
            text = response.choices[0].message.content or "{}"
            result = _parse_response_json(text, paths)
            if max_selections is not None and len(result) > max_selections:
                result = result[:max_selections]
            return result
        except Exception as e:
            last_error = e
            if attempt < OPENAI_MAX_RETRIES - 1:
                time.sleep(2 ** attempt)
            continue

    raise last_error or RuntimeError("Failed to get selection from OpenAI")