import os
import base64
import mimetypes
import logging
import asyncio
import io
import datetime
import time
from typing import Any, Callable, Dict, List, Union
from PIL import Image, ImageDraw, ImageFont
from tools.utils import format_relative_time

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 = args.get("path")
    page = int(args.get("page", 1))
    sort_by = args.get("sort_by", "mtime")
    
    if not dir_path:
        return {"text": "Error: Missing path argument"}
    if not os.path.isdir(dir_path):
        return {"text": f"Error: {dir_path} is not a directory."}
    
    exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
    file_info_list = []
    for f in os.listdir(dir_path):
        if f.lower().endswith(exts):
            full_path = os.path.join(dir_path, f)
            try:
                stats = os.stat(full_path)
                file_info_list.append({
                    "name": f,
                    "path": full_path,
                    "mtime": stats.st_mtime,
                    "size": stats.st_size
                })
            except Exception as e:
                logger.error(f"Could not stat {f}: {e}")

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

    # Sorting
    if sort_by == "mtime":
        file_info_list.sort(key=lambda x: x["mtime"], reverse=True)
        sort_label = "Date (newest first)"
    elif sort_by == "size":
        file_info_list.sort(key=lambda x: x["size"], reverse=True)
        sort_label = "Size (largest first)"
    else:
        file_info_list.sort(key=lambda x: x["name"].lower())
        sort_label = "Name (alphabetical)"

    # Fixed grid dimensions for token optimization
    COLS = 10
    ROWS = 7
    PAGE_SIZE = COLS * ROWS
    
    total_files = len(file_info_list)
    start_idx = (page - 1) * PAGE_SIZE
    end_idx = start_idx + PAGE_SIZE
    paged_files = file_info_list[start_idx:end_idx]
    
    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
    font_paths = [
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "/usr/share/fonts/truetype/liberation/LiberationSans-Bold.ttf",
        "/usr/share/fonts/truetype/freefont/FreeSansBold.ttf",
    ]
    for path in font_paths:
        if os.path.exists(path):
            try:
                font = ImageFont.truetype(path, 24)
                break
            except Exception:
                continue
    if font is None:
        font = ImageFont.load_default()

    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))
            
            # Index number in lower-left corner
            text = str(global_index)
            if font:
                bbox = draw.textbbox((0, 0), text, font=font)
                text_w = bbox[2] - bbox[0]
                text_h = bbox[3] - bbox[1]
            else:
                text_w = len(text) * 8
                text_h = 15
            
            # Semi-opaque background rectangle for the number
            draw.rectangle([x, y + thumb_size - text_h - 5, x + text_w + 5, y + thumb_size - 2], 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()
    canvas.save(buf, format='PNG')
    img_data = base64.b64encode(buf.getvalue()).decode("utf-8")
    
    total_pages = (total_files + PAGE_SIZE - 1) // PAGE_SIZE
    
    # Range shown: first and last of the CURRENT PAGE
    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}\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/png"}
    ]