import logging
from typing import Sequence, Union
from enum import Enum
import mcp.types as types
from mcp.server import Server
from gmail_mcp.utils.GCP.gmail_search import search_emails
from gmail_mcp.utils.GCP.gmail_auth import get_gmail_service

logger = logging.getLogger("gmail_mcp")


class ToolName(str, Enum):
    PERFORM_EMAIL_SEARCH = "perform-email-search"


ServerTools = [
    types.Tool(
        name=ToolName.PERFORM_EMAIL_SEARCH,
        description="Searches the user's Gmail with various filters and returns complete thread contents.",
        inputSchema={
            "type": "object",
            "properties": {
                "sender": {"type": "string"},
                "recipient": {"type": "string"},
                "subject": {"type": "string"},
                "has_attachment": {"type": "boolean"},
                "filename": {"type": "string"},
                "newer_than_days": {"type": "integer"},
                "older_than_days": {"type": "integer"},
                "label": {"type": "string"},
                "is_unread": {"type": "boolean"},
                "keywords": {"type": "string"},
            },
            "required": [],
        },
    ),
]


async def perform_email_search(
    sender: str = None,
    recipient: str = None,
    subject: str = None,
    has_attachment: bool = None,
    filename: str = None,
    days_after: int = None,
    days_before: int = None,
    newer_than_days: int = None,
    older_than_days: int = None,
    label: str = None,
    is_unread: bool = None,
    max_results: int = None,
    keywords: str = None,
    request_id: str = "",
    inflight_requests: dict = None
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Perform a Gmail search using the parameters provided, returning the combined thread contents.
    """

    if inflight_requests is not None:
        inflight_requests[request_id] = "running"

    logger.info("Initializing Gmail service for email search.")
    service = get_gmail_service()

    # Convert days_after / days_before into datetimes if needed
    from datetime import datetime, timedelta

    after_date = None
    if days_after is not None:
        after_date = datetime.now() - timedelta(days=days_after)

    before_date = None
    if days_before is not None:
        before_date = datetime.now() - timedelta(days=days_before)

    logger.info("Performing search with the provided query parameters.")
    results = search_emails(
        service=service,
        sender=sender,
        recipient=recipient,
        subject=subject,
        has_attachment=has_attachment,
        filename=filename,
        after_date=after_date,
        before_date=before_date,
        newer_than_days=newer_than_days,
        older_than_days=older_than_days,
        label=label,
        is_unread=is_unread,
        max_results=max_results,
        keywords=keywords
    )

    logger.info("Search completed. Returning formatted results.")

    if inflight_requests is not None:
        inflight_requests[request_id] = "done"

    if results.strip():
        return [types.TextContent(type="text", text=results)]
    else:
        return [types.TextContent(type="text", text="No results found.")]


def register_tools(server: Server, inflight_requests: dict):
    @server.list_tools()
    async def handle_list_tools() -> list[types.Tool]:
        return ServerTools

    @server.call_tool()
    async def handle_call_tool(
        name: str, arguments: dict | None
    ) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
        logger.info(f"Calling tool: {name} with arguments {arguments}")

        request_id = arguments.get("__request_id__", "no_request_id")
        if inflight_requests is None:
            logger.warning(
                "inflight_requests dictionary not provided, request not tracked.")

        if name == ToolName.PERFORM_EMAIL_SEARCH:
            return await perform_email_search(
                sender=arguments.get("sender"),
                recipient=arguments.get("recipient"),
                subject=arguments.get("subject"),
                has_attachment=arguments.get("has_attachment"),
                filename=arguments.get("filename"),
                days_after=arguments.get("days_after"),
                days_before=arguments.get("days_before"),
                newer_than_days=arguments.get("newer_than_days"),
                older_than_days=arguments.get("older_than_days"),
                label=arguments.get("label"),
                is_unread=arguments.get("is_unread"),
                max_results=arguments.get("max_results"),
                keywords=arguments.get("keywords"),
                request_id=request_id,
                inflight_requests=inflight_requests
            )
        else:
            logger.error(f"Unknown tool called: {name}")
            raise ValueError(f"Unknown tool: {name}")


__all__ = ["register_tools", "perform_email_search"]
