import base64
import logging
import io
import datetime
from pathlib import Path
from typing import Any, Dict, List
from PIL import Image, ImageDraw, ImageFont
from tools.utils import format_relative_time, ToolError, COMMON_FONTS, get_file_info_list, sort_file_list, get_paginated_list

logger = logging.getLogger("MattCP")

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 = 10
    ROWS = 7
    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
    for font_candidate in COMMON_FONTS:
        try:
            font = ImageFont.truetype(font_candidate, 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:
                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 quality 95 to save data
    canvas.save(buf, format='JPEG', quality=95)
    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"}
    ]