MooseCP/tools/contact_sheet.py

import base64
import logging
import io
import datetime
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.
    Grid: 10 columns x 7 rows (70 images total).
    Sized for optimal token usage (1120 tokens) on llama.cpp.
    """
    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)")

    # Fixed grid dimensions for token optimization
    COLS = config.CONTACT_SHEET_COLS
    ROWS = config.CONTACT_SHEET_ROWS
    PAGE_SIZE = COLS * 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}."}

    thumb_size = 192 # 192 / 48 = 4 patches per side
    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"}
    ]