from typing import Any, Callable, Dict

# Define the Tool class here so it's available to all handlers and the registry
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
        }

# Import handlers from separate files
from .read_image import handle as read_image_handler
from .read_metadata import handle as read_png_metadata_handler
from .contact_sheet import handle as contact_sheet_handler
from .list_directory import handle as list_directory_details_handler
from .preview_image import handle as preview_image_handler
from .get_text_context import handle as get_text_context_handler
from .browse_wikipedia import handle as wikipedia_handler
from .get_model_info import handle as get_model_info_handler
from .generate_image import handle as generate_image_handler
from .search_tags import handle as search_tags_handler

# Central registry of all available tools
TOOL_REGISTRY = [
    Tool(
        name="read_image",
        description="Reads an image at full resolution for maximum detail. Use this for final confirmation of a selected image or when deep visual analysis of a specific file is required.",
        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 prompts for AI generated images.",
        schema={
            "type": "object",
            "properties": {"path": {"type": "string", "description": "Path to the PNG file"}},
            "required": ["path"],
        },
        handler=read_png_metadata_handler
    ),
    Tool(
        name="contact_sheet",
        description="Generates a coarse, high-density overview of images in a directory. This tool is paginated; a single page may not show all images. To perform a comprehensive scan or find specific images, you MUST iterate through multiple pages. If ANYTHING requires selecting individual images, then use indices as inputs to preview_image. Rely on this tool ONLY for a general characterization of images in a directory, NOT anything about specific images.",
        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). Increment this value to navigate through the full list of images in the folder.", "default": 1},
                "sort_by": {"type": "string", "description": "Sort order: 'mtime' (newest first, default), 'name' (alphabetical), or 'size' (largest first)", "default": "mtime"},
            },
            "required": ["path"],
        },
        handler=contact_sheet_handler
    ),
    Tool(
        name="list_directory",
        description="Lists the contents of a directory in a simplified text format. Use this for checking file existence, sizes, or dates or for quick filesystem navigation. DO NOT use this tool for any sort of visual or image selection.",
        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},
                "sort_by": {"type": "string", "description": "Sort order: 'mtime' (newest first, default), 'name' (alphabetical), or 'size' (largest first)", "default": "mtime"},
            },
            "required": ["path"],
        },
        handler=list_directory_details_handler
    ),
    Tool(
        name="preview_image",
        description="Provides low-resolution previews and information such as filenames, size, and resolution for specific images. It uses the minimum viable token budget, so details WILL be missed. Only use this tool for quick down-selection of images from contact_sheet or where accuracy does not matter. For final confirmation or deep analysis of a single selected image, you MUST use read_image after this tool.",
        schema={
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path to an image file or a directory containing images"},
                "indices": {"type": "string", "description": "1-based indices or ranges (e.g., '1, 3-5, 10') to preview from the directory. Required if path is a directory."},
                "sort_by": {"type": "string", "description": "Sort order to resolve indices: 'mtime' (default), 'name', or 'size'."},
            },
            "required": ["path"],
        },
        handler=preview_image_handler
    ),
    Tool(
        name="get_text_context",
        description="Extracts specific sections of a file with absolute line numbers and surrounding context. This tool MUST be called immediately before edit_file to verify the exact line numbers and content of the block being replaced, preventing misalignment errors.",
        schema={
            "type": "object",
            "properties": {
                "path": {"type": "string", "description": "Path to the file"},
                "pattern": {"type": "string", "description": "Regex pattern to search for"},
                "lines": {"type": "string", "description": "1-based indices or ranges (e.g., '10-20, 45')"},
                "context": {"type": "integer", "description": "Number of lines of context to show above and below", "default": 1},
            },
            "required": ["path"],
        },
        handler=get_text_context_handler
    ),
    Tool(
        name="browse_wikipedia",
        description="Browse Wikipedia pages to retrieve information. Defaults to the page summary. Can also retrieve the table of contents or a specific section. If a page is not found, it automatically returns search results.",
        schema={
            "type": "object",
            "properties": {
                "title": {"type": "string", "description": "The title of the Wikipedia page to browse or the search query."},
                "mode": {
                    "type": "string", 
                    "description": "The retrieval mode. 'summary' (default) for the lead section, 'toc' for the table of contents, 'section' for a specific section, or 'search' for explicit search results.",
                    "enum": ["summary", "toc", "section", "search"],
                    "default": "summary"
                },
                "section_index": {"type": "integer", "description": "The linear index of the section to retrieve. Required if mode='section'. This index can be found by calling the tool in 'toc' mode."},
                "search_limit": {"type": "integer", "description": "The maximum number of search results to return. Default is 5.", "default": 5},
            },
            "required": ["title"],
        },
        handler=wikipedia_handler
    ),
    Tool(
        name="get_model_info",
        description="Returns a list of available txt2img models and their short descriptions. If a specific model name is provided, it returns the comprehensive prompting guide, available resolution presets, and active configuration tips for that model. CRITICAL: You must call this for any model you are unfamiliar with, as prompting styles vary wildly (e.g., tag-based vs. natural language) and using the wrong style will result in poor image quality.",
        schema={
            "type": "object",
            "properties": {
                "model_name": {"type": "string", "description": "The name of the model to get detailed info for. Leave empty to list all available models."}
            },
            "required": [],
        },
        handler=get_model_info_handler
    ),
    Tool(
        name="generate_image",
        description="Generates an image using the specified model and parameters. To ensure high quality, verify the model's prompting requirements via get_model_info before calling this tool. You will usually not get the correct result the first time. If the generated image needs work, change the prompt and call this tool again. If it's good, use the provided markdown to display it to the user.",
        schema={
            "type": "object",
            "properties": {
                "model_name": {"type": "string", "description": "The name of the model to use. Required."},
                "prompt": {"type": "string", "description": "The prompt for the image. Required."},
                "negative_prompt": {"type": "string", "description": "The negative prompt to exclude unwanted elements."},
                "resolution_preset": {"type": "string", "description": "A named resolution preset from the model info"},
                "cfg_scale": {"type": "number", "description": "CFG scale for prompt adherence. Usually should be left omitted to select the default."},
            },
            "required": ["model_name", "prompt", "resolution_preset"],
        },
        handler=generate_image_handler
    ),
    Tool(
        name="search_tags",
        description="Searches the Danbooru tag database for tags matching a query. Returns the Danbooru wiki page on an exact match. On any query, returns a list of similar tags as well as aliases. Useful for finding the correct booru-style tags for anime models or getting more information.",
        schema={
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "The tag or partial tag to search for (e.g., 'white hair' or 'lesbian')."}
            },
            "required": ["query"],
        },
        handler=search_tags_handler
    ),
]
