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
from tools.utils import format_relative_time, ToolError, get_file_info_list, sort_file_list

logger = logging.getLogger("MattCP")

def parse_indices(indices_str: str) -> List[int]:
    """
    Parses a string like '11, 31-33, 44' into a list of 1-based indices.
    Preserves the order specified in the string.
    """
    indices = []
    parts = [p.strip() for p in indices_str.split(',')]
    for part in parts:
        if not part:
            continue
        if '-' in part:
            try:
                start, end = map(int, part.split('-'))
                # Expand range in order. Handle reverse ranges as well.
                step = 1 if start <= end else -1
                for i in range(start, end + step, step):
                    indices.append(i)
            except ValueError:
                raise ToolError(f"Invalid range format: {part}")
        else:
            try:
                indices.append(int(part))
            except ValueError:
                raise ToolError(f"Invalid index format: {part}")
    return indices

def calculate_patch_dimensions(orig_w: int, orig_h: int, target_tokens: int = 70):
    """
    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.
    """
    ar = orig_w / orig_h
    
    # Case A: Width is the anchor (non-padded)
    w_anchor = 1
    while True:
        w_px = w_anchor * 48
        h_px = round(w_px / ar)
        # Calculate actual tokens based on resulting pixel dimensions
        tokens = math.ceil(w_px / 48) * math.ceil(h_px / 48)
        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 * 48
        w_px = round(h_px * ar)
        # Calculate actual tokens based on resulting pixel dimensions
        tokens = math.ceil(w_px / 48) * math.ceil(h_px / 48)
        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=95)
                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