import math
import base64
import logging
import io
import datetime
from pathlib import Path
from typing import Any, Dict, List, Tuple
from PIL import Image, ImageOps
import config
from tools.utils import format_relative_time, ToolError, get_file_info_list, sort_file_list, parse_indices


logger = logging.getLogger("MooseCP")

def calculate_patch_dimensions(orig_w: int, orig_h: int, target_tokens: int = None):
    """
    Calculates optimal pixel dimensions to hit a token budget of at least target_tokens.
    Validates budget based on actual rounded pixel dimensions to avoid float precision errors.
    """
    if target_tokens is None:
        target_tokens = config.PREVIEW_TOKEN_BUDGET
    
    ar = orig_w / orig_h
    
    # Case A: Width is the anchor (non-padded)
    w_anchor = 1
    while True:
        w_px = w_anchor * config.PATCH_SIZE
        h_px = round(w_px / ar)
        # Calculate actual tokens based on resulting pixel dimensions
        tokens = math.ceil(w_px / config.PATCH_SIZE) * math.ceil(h_px / config.PATCH_SIZE)
        if tokens >= target_tokens:
            res_a = {"w": w_px, "h": h_px, "tokens": tokens}
            break
        w_anchor += 1

    # Case B: Height is the anchor (non-padded)
    h_anchor = 1
    while True:
        h_px = h_anchor * config.PATCH_SIZE
        w_px = round(h_px * ar)
        # Calculate actual tokens based on resulting pixel dimensions
        tokens = math.ceil(w_px / config.PATCH_SIZE) * math.ceil(h_px / config.PATCH_SIZE)
        if tokens >= target_tokens:
            res_b = {"w": w_px, "h": h_px, "tokens": tokens}
            break
        h_anchor += 1

    # Pick the one with the smaller token count
    best = res_a if res_a["tokens"] <= res_b["tokens"] else res_b
    return max(1, best["w"]), max(1, best["h"])

async def handle(args: Dict[str, Any]):
    """
    Generates small thumbnails of images with detailed info.
    Can take a single file path, or a directory with indices/ranges from contact_sheet.
    Thumbnails are optimized for ~70 token usage.
    """
    path_str = args.get("path")
    indices_str = args.get("indices")
    sort_by = args.get("sort_by", "mtime")
    
    if not path_str:
        raise ToolError("Missing path argument")
    
    path = Path(path_str)
    # Store as (path, index_label) where index_label is None if direct path
    target_files: List[Tuple[Path, Any]] = []

    if path.is_file():
        target_files.append((path, None))
    elif path.is_dir():
        if not indices_str:
            raise ToolError("Indices are required when providing a directory path. Example: '1, 5-10'")
        
        exts = ('.png', '.jpg', '.jpeg', '.webp', '.bmp')
        file_info_list = get_file_info_list(path, extensions=exts)
        file_info_list = sort_file_list(file_info_list, sort_by)
        
        indices = parse_indices(indices_str)
        for idx in indices:
            # contact_sheet uses 1-based indexing
            if 1 <= idx <= len(file_info_list):
                target_files.append((file_info_list[idx-1]["path"], idx))
            else:
                logger.warning(f"Index {idx} out of range for directory {path_str}")
    else:
        raise ToolError(f"Path {path_str} is neither a file nor a directory.")

    if not target_files:
        return {"text": "No valid images found to preview."}

    results = []
    for file_path, index_label in target_files:
        try:
            stats = file_path.stat()
            with Image.open(file_path) as img:
                # Apply EXIF orientation to fix sideways images
                img = ImageOps.exif_transpose(img)
                orig_w, orig_h = img.size
                target_w, target_h = calculate_patch_dimensions(orig_w, orig_h)
                
                # Use thumbnail() to preserve aspect ratio and fit within the target bounding box.
                thumb = img.convert("RGB")
                thumb.thumbnail((target_w, target_h), Image.Resampling.LANCZOS)
                
                buf = io.BytesIO()
                thumb.save(buf, format='JPEG', quality=config.IMAGE_QUALITY)
                img_data = base64.b64encode(buf.getvalue()).decode("utf-8")
                
                # Gather info
                mtime_rel = format_relative_time(stats.st_mtime)
                mtime_abs = datetime.datetime.fromtimestamp(stats.st_mtime).strftime('%Y-%m-%d %H:%M')
                size_kb = stats.st_size / 1024
                
                # Prepend index if this image was selected via index/range
                idx_prefix = f"[Index: {index_label}] " if index_label is not None else ""
                
                info_text = (
                    f"{idx_prefix}File: {file_path.name}\n"
                    f"Resolution: {orig_w}x{orig_h}\n"
                    f"Size: {size_kb:.1f}KB\n"
                    f"Modified: {mtime_rel} ({mtime_abs})"
                )
                
                results.append({"type": "text", "text": info_text})
                results.append({"type": "image", "data": img_data, "mimeType": "image/jpeg"})
                
        except Exception as e:
            results.append({"type": "text", "text": f"Error processing {file_path.name}: {e}"})

    return results