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

logger = logging.getLogger("MattCP")

def format_relative_time(timestamp: float) -> str:
    """Converts a timestamp to a human-readable relative format."""
    now = time.time()
    diff = now - timestamp
    seconds = int(diff)
    if seconds < 60:
        return "Just now"
    minutes = seconds // 60
    if minutes < 60:
        return f"{minutes}m ago"
    hours = minutes // 60
    if hours < 24:
        return f"{hours}h ago"
    days = hours // 24
    if days == 1:
        return "Yesterday"
    if days < 7:
        return f"{days}d ago"
    weeks = days // 7
    if weeks < 4:
        return f"{weeks}w ago"
    dt = datetime.datetime.fromtimestamp(timestamp)
    return dt.strftime('%Y-%m-%d')

class Tool:
    def __init__(self, name: str, description: str, schema: Dict[str, Any], handler: Callable):
        self.name = name
        self.description = description
        self.schema = schema
        self.handler = handler

    def to_mcp_dict(self):
        return {
            "name": self.name,
            "description": self.description,
            "inputSchema": self.schema
        }

async def read_image_handler(args: Dict[str, Any]):
    path = args.get("path")
    if not path:
        return {"text": "Error: Missing path argument"}
    mime_type, _ = mimetypes.guess_type(path)
    if not mime_type or not mime_type.startswith("image/"):
        return {"text": f"Error: File {path} is not a supported image format."}
    try:
        with open(path, "rb") as f:
            data = base64.b64encode(f.read()).decode("utf-8")
            return {"type": "image", "data": data, "mimeType": mime_type}
    except Exception as e:
        return {"text": f"Error reading image: {str(e)}"}

async def read_png_metadata_handler(args: Dict[str, Any]):
    path = args.get("path")
    if not path:
        return {"text": "Error: Missing path argument"}
    try:
        with Image.open(path) as img:
            if img.format != 'PNG':
                return {"text": "Error: File is not a PNG image."}
            metadata = img.info
            if not metadata:
                return {"text": "No metadata found in this PNG."}
            output = "\n".join([f"{k}: {v}" for k, v in metadata.items()])
            return {"text": f"PNG Metadata:\n{output}"}
    except Exception as e:
        return {"text": f"Error reading PNG metadata: {str(e)}"}

async def list_thumbnails_handler(args: Dict[str, Any]):
    dir_path = args.get("path")
    page = int(args.get("page", 1))
    page_size = int(args.get("page_size", 64))
    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."}

    if sort_by == "mtime":
        file_info_list.sort(key=lambda x: x["mtime"], reverse=True)
    elif sort_by == "size":
        file_info_list.sort(key=lambda x: x["size"], reverse=True)
    else:
        file_info_list.sort(key=lambda x: x["name"].lower())

    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 = 160 
    padding = 15
    label_height = 35 
    
    num_files = len(paged_files)
    cols = int(num_files**0.5) if num_files > 0 else 1
    if cols == 0: cols = 1
    rows = (num_files + cols - 1) // cols
    
    canvas_w = cols * (thumb_size + padding) + padding
    canvas_h = rows * (thumb_size + label_height + padding) + padding
    
    canvas = Image.new('RGB', (canvas_w, canvas_h), (30, 30, 30))
    draw = ImageDraw.Draw(canvas)
    
    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, 20)
                break
            except Exception:
                continue
    if font is None:
        font = ImageFont.load_default()

    file_details = []

    for i, info in enumerate(paged_files):
        filename = info["name"]
        full_path = info["path"]
        row = i // cols
        col = i % cols
        
        x = padding + col * (thumb_size + padding)
        y = padding + row * (thumb_size + label_height + padding)
        
        # Continuous numbering across pages
        global_index = start_idx + i + 1
        
        try:
            with Image.open(full_path) as img:
                width, height = img.size
                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))
            
            mod_time_rel = format_relative_time(info["mtime"])
            size_kb = info["size"] / 1024
            file_details.append(f"{global_index}. {filename} | {width}x{height} | {size_kb:.1f}KB | {mod_time_rel}")
            
            text = str(global_index)
            draw.text((x + 2, y + thumb_size + 2), text, fill=(255, 255, 255), font=font)
            
        except Exception as e:
            logger.error(f"Failed to process {filename}: {e}")
            draw.text((x, y + thumb_size // 2), "Error", fill=(255, 0, 0), font=font)
            file_details.append(f"{global_index}. {filename} | ERROR")

    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
    header_text = f"Directory: {dir_path}\nPage {page} of {total_pages} ({total_files} total images). Showing {start_idx+1}-{min(end_idx, total_files)}."
    
    details_text = "\n".join(file_details)
    
    return [
        {"type": "text", "text": f"{header_text}\n\nFile List:\n{details_text}"},
        {"type": "image", "data": img_data, "mimeType": "image/png"}
    ]

async def list_directory_details_handler(args: Dict[str, Any]):
    path = args.get("path")
    page = int(args.get("page", 1))
    page_size = int(args.get("page_size", 64))
    
    if not path:
        return {"text": "Error: Missing path argument"}
    if not os.path.isdir(path):
        return {"text": f"Error: {path} is not a directory."}

    try:
        entries = os.listdir(path)
    except Exception as e:
        return {"text": f"Error reading directory: {e}"}

    all_items = []
    for entry in entries:
        full_path = os.path.join(path, entry)
        try:
            stats = os.stat(full_path)
            mtime_rel = format_relative_time(stats.st_mtime)
            
            if os.path.isdir(full_path):
                try:
                    count = len(os.listdir(full_path))
                except:
                    count = 0
                all_items.append({
                    "type": "dir",
                    "name": entry,
                    "info": f"{count} items",
                    "mtime": mtime_rel
                })
            else:
                size_kb = stats.st_size / 1024
                all_items.append({
                    "type": "file",
                    "name": entry,
                    "info": f"{size_kb:.1f}KB",
                    "mtime": mtime_rel
                })
        except Exception as e:
            logger.error(f"Could not stat {entry}: {e}")

    all_items.sort(key=lambda x: (x["type"] == "file", x["name"].lower()))

    total_items = len(all_items)
    start_idx = (page - 1) * page_size
    end_idx = start_idx + page_size
    paged_items = all_items[start_idx:end_idx]

    if not paged_items:
        return {"text": f"No items found on page {page}."}

    output = []
    for item in paged_items:
        if item["type"] == "dir":
            output.append(f"[DIR] {item['name']} | {item['info']} | {item['mtime']}")
        else:
            output.append(f"      {item['name']} | {item['info']} | {item['mtime']}")

    total_pages = (total_items + page_size - 1) // page_size
    header = f"Listing of {path}\nPage {page} of {total_pages} ({total_items} total items). Showing {start_idx+1}-{min(end_idx, total_items)}."
    
    return {"text": f"{header}\n\n" + "\n".join(output)}

TOOL_REGISTRY = [
    Tool(
        name="read_image",
        description="Reads an image from the disk and returns it as an image content object.",
        schema={
            "type": "object",
            "properties": {"path": {"type": "string", "description": "Path to the image file"}},
            "required": ["path"],
        },
        handler=read_image_handler
    ),
    Tool(
        name="read_png_metadata",
        description="Reads the metadata (tEXt, zTXt, iTXt) from a PNG image to fetch info such as AI prompts.",
        schema={
            "type": "object",
            "properties": {"path": {"type": "string", "description": "Path to the PNG file"}},
            "required": ["path"],
        },
        handler=read_png_metadata_handler
    ),
    Tool(
        name="list_thumbnails",
        description="Generates a thumbnail grid of images in a directory. Supports pagination and sorting.",
        schema={
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path to the directory containing images"},
                "page": {"type": "integer", "description": "The page number to retrieve (1-indexed)", "default": 1},
                "page_size": {"type": "integer", "description": "Number of images per page", "default": 64},
                "sort_by": {"type": "string", "description": "Sort order: 'mtime' (newest first, default), 'name' (alphabetical), or 'size' (largest first)", "default": "mtime"},
            },
            "required": ["path"],
        },
        handler=list_thumbnails_handler
    ),
    Tool(
        name="list_directory",
        description="Lists the contents of a directory in a file browser format. Directories first, then files. Supports pagination.",
        schema={
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path to the directory to list"},
                "page": {"type": "integer", "description": "The page number to retrieve (1-indexed)", "default": 1},
                "page_size": {"type": "integer", "description": "Number of items per page", "default": 64},
            },
            "required": ["path"],
        },
        handler=list_directory_details_handler
    ),
]