MooseCP/tools/contact_sheet.py

import base64
import logging
import io
import datetime
import math
from pathlib import Path

from typing import Any, Dict, List
from PIL import Image, ImageDraw, ImageFont, ImageOps
import config
from tools.utils import format_relative_time, ToolError, get_file_info_list, sort_file_list, get_paginated_list


logger = logging.getLogger("MooseCP")

async def handle(args: Dict[str, Any]):
    """
    Generates a contact sheet of images.
    Sized for optimal token usage according to config values.
    """
    dir_path_str = args.get("path")
    page = int(args.get("page", 1))
    sort_by = args.get("sort_by", "mtime")
    
    if not dir_path_str:
        raise ToolError("Missing path argument")
    
    dir_path = Path(dir_path_str)
    if not dir_path.is_dir():
        raise ToolError(f"{dir_path_str} is not a directory.")
    
    exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
    file_info_list = get_file_info_list(dir_path, extensions=exts)

    if not file_info_list:
        return {"text": "No supported images found in the directory."}

    # Sorting
    file_info_list = sort_file_list(file_info_list, sort_by)
    
    sort_labels = {
        "mtime": "Date (newest first)",
        "size": "Size (largest first)",
        "name": "Name (alphabetical)"
    }
    sort_label = sort_labels.get(sort_by, "Name (alphabetical)")

    # Grid dimensions for token optimization (see config.py)
    MAX_COLS = config.CONTACT_SHEET_COLS
    MAX_ROWS = config.CONTACT_SHEET_ROWS
    PAGE_SIZE = MAX_COLS * MAX_ROWS

    
    total_files = len(file_info_list)
    paged_files = get_paginated_list(file_info_list, page, PAGE_SIZE)
    
    if not paged_files:
        return {"text": f"No images found on page {page}."}

    # Dynamic grid sizing to avoid blank space
    n_images = len(paged_files)
    thumb_size = config.CONTACT_SHEET_THUMB_SIZE
    
    if n_images <= config.CONTACT_SHEET_ROWS ** 2:
        # Smallest square-ish grid that fits all
        cols = math.ceil(math.sqrt(n_images))
        rows = math.ceil(n_images / cols)
    else:
        # Lock rows, expand columns
        rows = config.CONTACT_SHEET_ROWS
        cols = math.ceil(n_images / config.CONTACT_SHEET_ROWS)
    
    canvas_w = cols * thumb_size
    canvas_h = rows * thumb_size
    
    canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30))
    draw = ImageDraw.Draw(canvas)
    
    # Font loading
    font = None
    # Try generic system font names first
    for font_name in config.SYSTEM_FONT_NAMES:
        try:
            font = ImageFont.truetype(font_name, 24)
            break
        except Exception:
            continue
    # Fallback to absolute paths
    if font is None:
        for font_path in config.FALLBACK_FONT_PATHS:
            try:
                font = ImageFont.truetype(font_path, 24)
                break
            except Exception:
                continue
    if font is None:
        font = ImageFont.load_default()


    start_idx = (page - 1) * PAGE_SIZE
    for i, info in enumerate(paged_files):
        full_path = info["path"]
        row = i // cols
        col = i % cols
        
        x = col * thumb_size
        y = row * thumb_size
        global_index = start_idx + i + 1
        
        try:
            with Image.open(full_path) as img:
                # Apply EXIF orientation to fix sideways images
                img = ImageOps.exif_transpose(img)
                img.thumbnail((thumb_size, thumb_size))
                off_x = (thumb_size - img.width) // 2
                off_y = (thumb_size - img.height) // 2
                canvas.paste(img, (x + off_x, y + off_y))
            
            text = str(global_index)
            if font:
                bbox = draw.textbbox((0, 0), text, font=font)
                text_w, text_h = bbox[2] - bbox[0], bbox[3] - bbox[1]
            else:
                text_w, text_h = len(text) * 8, 15
            
            # Adjusted black box to align perfectly with the bottom of the image (y + thumb_size)
            draw.rectangle([x, y + thumb_size - text_h - 4, x + text_w + 5, y + thumb_size], fill=(0, 0, 0, 180))
            draw.text((x + 2, y + thumb_size - text_h - 7), text, fill=(255, 255, 255), font=font)
            
        except Exception as e:
            logger.error(f"Failed to process {info['name']}: {e}")
            draw.text((x + 5, y + thumb_size // 2), "Error", fill=(255, 0, 0), font=font)

    buf = io.BytesIO()
    # Save as JPEG to save data
    canvas.save(buf, format='JPEG', quality=config.IMAGE_QUALITY)
    img_data = base64.b64encode(buf.getvalue()).decode("utf-8")
    
    total_pages = (total_files + PAGE_SIZE - 1) // PAGE_SIZE
    
    first_on_page = paged_files[0]
    last_on_page = paged_files[-1]
    
    if sort_by == "mtime":
        first_val = f"{format_relative_time(first_on_page['mtime'])} ({datetime.datetime.fromtimestamp(first_on_page['mtime']).strftime('%Y-%m-%d')})"
        last_val = f"{format_relative_time(last_on_page['mtime'])} ({datetime.datetime.fromtimestamp(last_on_page['mtime']).strftime('%Y-%m-%d')})"
    elif sort_by == "size":
        first_val = f"{first_on_page['size']/1024:.1f}KB"
        last_val = f"{last_on_page['size']/1024:.1f}KB"
    else:
        first_val = first_on_page['name']
        last_val = last_on_page['name']
        
    header_text = (
        f"Directory: {dir_path_str}\n"
        f"Page {page} of {total_pages} ({total_files} images total). Sorted by: {sort_label}\n"
        f"Range shown: {first_val} ... {last_val}"
    )
    
    return [
        {"type": "text", "text": header_text},
        {"type": "image", "data": img_data, "mimeType": "image/jpeg"}
    ]