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
import aiohttp
from gmail_mcp.utils.openrouter.openrouter_utils import make_api_call, modify_draft_prompt
from gmail_mcp.utils.GCP.gmail_indexer import GmailIndexer
import base64

logger = logging.getLogger("gmail_mcp")


class ToolName(str, Enum):
    PERFORM_EMAIL_SEARCH = "perform-email-search"
    MODIFY_DRAFT_PROMPT = "modify-draft-prompt"
    LIST_NEW_EMAILS = "list-new-emails"
    GET_DRAFT_PROMPT = "get-draft-prompt"
    WRITE_DRAFT_EMAIL = "write-draft-email"
    SEND_APPROVED_EMAIL = "send-approved-email"


ServerTools = [
    types.Tool(
        name=ToolName.PERFORM_EMAIL_SEARCH,
        description="Searches the user's Gmail with various filters and returns complete contents.",
        inputSchema={
            "type": "object",
            "properties": {
                "sender": {"type": "string"},
                "recipient": {"type": "string"},
                "subject": {"type": "string"},
                "has_attachment": {"type": "boolean"},
                "filedoes ame": {"type": "string"},
                "newer_than_days": {"type": "integer"},
                "older_than_days": {"type": "integer"},
                "label": {"type": "string"},
                "is_unread": {"type": "boolean"},
                "keywords": {"type": "string"},
            },
            "required": [],
        },
    ),
    types.Tool(
        name=ToolName.MODIFY_DRAFT_PROMPT,
        description="Automatically adapts email writing style when user expresses preferences or concerns about communication style. This tool helps maintain personalized communication patterns for different contacts and contexts.",
        inputSchema={
            "type": "object",
            "properties": {
                "new_requirements": {
                    "type": "string",
                    "description": "Detailed context about the desired writing style changes, including specific scenarios, recipient relationships, and communication preferences. The more context provided (e.g., 'when emailing Professor Smith, I prefer...', 'for client communications, we should...'), the better the adaptation will be."
                }
            },
            "required": ["new_requirements"],
        },
    ),
    types.Tool(
        name=ToolName.LIST_NEW_EMAILS,
        description="Retrieves and lists all new emails found during Gmail index update.",
        inputSchema={
            "type": "object",
            "properties": {},
            "required": [],
        },
    ),
    types.Tool(
        name=ToolName.GET_DRAFT_PROMPT,
        description="""Retrieves the draft prompt template for composing a response.
        
        IMPORTANT: After getting this prompt:
        1. Use it IMMEDIATELY to write your draft
        2. Call write-draft-email with your draft
        
        Only call this tool when you have all needed information and are ready to write.""",
        inputSchema={
            "type": "object",
            "properties": {},
            "required": [],
        },
    ),
    types.Tool(
        name=ToolName.WRITE_DRAFT_EMAIL,
        description="""Shows the draft email to the user in the chat.
        
        IMPORTANT: After calling this tool:
        1. Present the draft content directly in your message (don't leave it in the expandable window)
        2. Ask the user to review and reply with 'SEND' to approve
        
        This tool helps ensure drafts are always properly shown to users before sending.""",
        inputSchema={
            "type": "object",
            "properties": {
                "to": {
                    "type": "string",
                    "description": "Recipient email address"
                },
                "subject": {
                    "type": "string",
                    "description": "Email subject"
                },
                "message_id": {
                    "type": "string",
                    "description": "Original message ID being replied to"
                },
                "thread_id": {
                    "type": "string",
                    "description": "Thread ID"
                },
                "draft_content": {
                    "type": "string",
                    "description": "The draft email content to be shown to the user"
                }
            },
            "required": ["to", "subject", "message_id", "thread_id", "draft_content"],
        },
    ),
    types.Tool(
        name=ToolName.SEND_APPROVED_EMAIL,
        description="""Sends an email that was previously shown and approved by the user.
        
        IMPORTANT: This tool must ONLY be called after:
        1. A draft was shown using write-draft-email
        2. The user responded with exactly 'SEND'
        
        NEVER call this tool without first showing the draft via write-draft-email.""",
        inputSchema={
            "type": "object",
            "properties": {
                "to": {
                    "type": "string",
                    "description": "Recipient email address"
                },
                "subject": {
                    "type": "string",
                    "description": "Email subject"
                },
                "message_id": {
                    "type": "string",
                    "description": "Original message ID being replied to"
                },
                "thread_id": {
                    "type": "string",
                    "description": "Thread ID"
                },
                "content": {
                    "type": "string",
                    "description": "The exact draft content that was shown and approved"
                },
                "user_confirmation": {
                    "type": "string",
                    "description": "Must be exactly 'SEND' in response to the draft",
                    "enum": ["SEND"]
                }
            },
            "required": ["to", "subject", "message_id", "thread_id", "content", "user_confirmation"],
        },
    ),
]


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.")]


async def modify_draft_prompt_tool(
    new_requirements: str,
    request_id: str = "",
    inflight_requests: dict = None
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Modify the existing draft prompt with new requirements.
    """
    if inflight_requests is not None:
        inflight_requests[request_id] = "running"

    try:
        # Read the current prompt
        with open("src/gmail_mcp/utils/constants/draft_prompt.txt", "r") as f:
            current_prompt = f.read()

        # Modify the prompt
        async with aiohttp.ClientSession() as session:
            modified_prompt = await modify_draft_prompt(current_prompt, new_requirements, session)

        # Save the modified prompt
        with open("src/gmail_mcp/utils/constants/draft_prompt.txt", "w") as f:
            f.write(modified_prompt)

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

        return [types.TextContent(type="text", text="Draft prompt successfully modified.")]

    except Exception as e:
        logger.error(f"Error in modify_draft_prompt_tool: {e}")
        if inflight_requests is not None:
            inflight_requests[request_id] = "error"
        return [types.TextContent(type="text", text=f"Error modifying draft prompt: {str(e)}")]


async def list_new_emails(
    request_id: str = "",
    inflight_requests: dict = None
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Retrieve and list only the new emails found during index update.
    """
    if inflight_requests is not None:
        inflight_requests[request_id] = "running"

    try:
        service = get_gmail_service()
        indexer = GmailIndexer(service)

        # Get new emails using the dedicated method
        new_emails = indexer.get_new_emails()

        if not new_emails:
            if inflight_requests is not None:
                inflight_requests[request_id] = "done"
            return [types.TextContent(type="text", text="No new emails found.")]

        output = f"Found {len(new_emails)} new emails:\n\n"
        for email in new_emails:
            output += f"Subject: {email.get('subject', 'No subject')}\n"
            output += f"From: {email.get('from', 'Unknown sender')}\n"
            output += f"Date: {email.get('date', 'No date')}\n"
            output += f"Snippet: {email.get('snippet', 'No preview available')}\n"
            output += f"Thread ID: {email.get('threadId', 'No thread ID')}\n"
            output += f"Message ID: {email.get('id', 'No message ID')}\n"
            output += "-" * 50 + "\n"

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

        return [types.TextContent(type="text", text=output)]

    except Exception as e:
        logger.error(f"Error in list_new_emails: {e}")
        if inflight_requests is not None:
            inflight_requests[request_id] = "error"
        return [types.TextContent(type="text", text=f"Error listing new emails: {str(e)}")]


async def get_draft_prompt(
    request_id: str = "",
    inflight_requests: dict = None
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Get the draft prompt template for composing a response.
    """
    if inflight_requests is not None:
        inflight_requests[request_id] = "running"

    try:
        # Get the current draft prompt
        with open("src/gmail_mcp/utils/constants/draft_prompt.txt", "r") as f:
            draft_prompt = f.read()

        # Append fixed instructions for information gathering
        fixed_instructions = """

INFORMATION GATHERING INSTRUCTIONS:

1. For User-Specific Information:
   - Search Gmail history first using the search_emails tool
   - Do not search company database for personal information
   - If no relevant information is found, ask the user specific questions

2. For Company-Related Information:
   - Search both Gmail history AND company database
   - Use search_emails tool for Gmail
   - Use perform_company_search tool for company database
   - Only after exhausting both sources, request user input

3. When Requesting User Input:
   - Break down complex questions into simple, direct questions
   - Format each question clearly and separately
   - Specify exactly what information is needed to complete the response
   - Example:
     Instead of: "What's your availability and preferences for the project timeline?"
     Ask:
     1. What days are you available next week?
     2. What is your preferred meeting time (morning/afternoon)?
     3. How many hours can you commit to this project weekly?

Remember: Always explain what sources were checked before requesting user input."""

        complete_prompt = draft_prompt + fixed_instructions

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

        return [types.TextContent(type="text", text=complete_prompt)]

    except Exception as e:
        logger.error(f"Error in get_draft_prompt: {e}")
        if inflight_requests is not None:
            inflight_requests[request_id] = "error"
        return [types.TextContent(type="text", text=f"Error getting draft prompt: {str(e)}")]


async def send_approved_email(
    to: str,
    subject: str,
    message_id: str,
    thread_id: str,
    content: str,
    user_confirmation: str,
    request_id: str = "",
    inflight_requests: dict = None
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Sends an email that has been explicitly approved by the user.
    This should only be called after the draft has been shown and approved in the chat.
    """
    if user_confirmation != "SEND":
        return [types.TextContent(type="text", text="Email not sent. User must explicitly type 'SEND' to approve sending.")]

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

    try:
        service = get_gmail_service()

        # Clean up the IDs
        thread_id = thread_id.replace('thread_', '').replace(
            'new_thread_', '').replace('new_', '')
        message_id = message_id.replace('msg_', '').replace('message_', '')

        # Create MIME message
        from email.mime.text import MIMEText
        from email.mime.multipart import MIMEMultipart

        message = MIMEMultipart()
        message['to'] = to
        message['subject'] = subject
        message['In-Reply-To'] = f"<{message_id}@mail.gmail.com>"
        message['References'] = f"<{message_id}@mail.gmail.com>"

        # Add the message body
        msg = MIMEText(content, 'plain')
        message.attach(msg)

        # Create and send the message
        email = {
            'raw': base64.urlsafe_b64encode(message.as_bytes()).decode(),
            'threadId': thread_id,
            'labelIds': ['SENT']
        }

        sent_message = service.users().messages().send(userId='me', body=email).execute()

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

        return [types.TextContent(type="text", text=f"Email sent successfully. Message ID: {sent_message['id']}")]

    except Exception as e:
        logger.error(f"Error in send_approved_email: {e}")
        if inflight_requests is not None:
            inflight_requests[request_id] = "error"
        return [types.TextContent(type="text", text=f"Error sending email: {str(e)}")]


async def write_draft_email(
    to: str,
    subject: str,
    message_id: str,
    thread_id: str,
    draft_content: str,
    request_id: str = "",
    inflight_requests: dict = None
) -> Sequence[Union[types.TextContent, types.ImageContent, types.EmbeddedResource]]:
    """
    Returns the draft with instructions for the AI to present it to the user.
    """
    if inflight_requests is not None:
        inflight_requests[request_id] = "running"

    try:
        response = f"""Here is the draft email. Rewrite it to the user in the chat again and ask them to reply with 'SEND' to approve:

To: {to}
Subject: {subject}

{draft_content}"""

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

        return [types.TextContent(type="text", text=response)]

    except Exception as e:
        logger.error(f"Error in write_draft_email: {e}")
        if inflight_requests is not None:
            inflight_requests[request_id] = "error"
        return [types.TextContent(type="text", text=f"Error preparing draft email: {str(e)}")]


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
            )
        elif name == ToolName.MODIFY_DRAFT_PROMPT:
            return await modify_draft_prompt_tool(
                new_requirements=arguments.get("new_requirements"),
                request_id=request_id,
                inflight_requests=inflight_requests
            )
        elif name == ToolName.LIST_NEW_EMAILS:
            return await list_new_emails(
                request_id=request_id,
                inflight_requests=inflight_requests
            )
        elif name == ToolName.GET_DRAFT_PROMPT:
            return await get_draft_prompt(
                request_id=request_id,
                inflight_requests=inflight_requests
            )
        elif name == ToolName.WRITE_DRAFT_EMAIL:
            return await write_draft_email(
                to=arguments.get("to"),
                subject=arguments.get("subject"),
                message_id=arguments.get("message_id"),
                thread_id=arguments.get("thread_id"),
                draft_content=arguments.get("draft_content"),
                request_id=request_id,
                inflight_requests=inflight_requests
            )
        elif name == ToolName.SEND_APPROVED_EMAIL:
            return await send_approved_email(
                to=arguments.get("to"),
                subject=arguments.get("subject"),
                message_id=arguments.get("message_id"),
                thread_id=arguments.get("thread_id"),
                content=arguments.get("content"),
                user_confirmation=arguments.get("user_confirmation"),
                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"]
