from typing import Optional, List
from datetime import datetime
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
import base64
import logging

logger = logging.getLogger("gmail_mcp")


def search_emails(
    service,
    sender: Optional[str] = None,
    recipient: Optional[str] = None,
    subject: Optional[str] = None,
    has_attachment: Optional[bool] = None,
    filename: Optional[str] = None,
    after_date: Optional[datetime] = None,
    before_date: Optional[datetime] = None,
    newer_than_days: Optional[int] = None,
    older_than_days: Optional[int] = None,
    label: Optional[str] = None,
    is_unread: Optional[bool] = None,
    max_results: Optional[int] = None,
    keywords: Optional[str] = None
) -> str:
    """
    Search Gmail with various filters and return complete thread contents.

    Args:
        service: Gmail API service instance
        sender: Email address of the sender
        recipient: Email address of the recipient
        subject: Subject line content to search for
        has_attachment: If True, only returns messages with attachments
        filename: Search for specific attachment filename or extension
        after_date: Return messages after this date (datetime object)
        before_date: Return messages before this date (datetime object)
        newer_than_days: Messages newer than specified days
        older_than_days: Messages older than specified days
        label: Gmail label to search within
        is_unread: If True, only returns unread messages
        max_results: Maximum number of results to return
        keywords: Text to search for in email body or subject (space-separated for multiple terms)

    Returns:
        str: A formatted string containing all thread contents
    """
    # Add query construction logging
    query_parts = []
    logger.info("Building search query with parameters:")
    logger.info(f"  - sender: {sender}")
    logger.info(f"  - recipient: {recipient}")
    logger.info(f"  - subject: {subject}")
    logger.info(f"  - has_attachment: {has_attachment}")
    logger.info(f"  - newer_than_days: {newer_than_days}")
    logger.info(f"  - older_than_days: {older_than_days}")
    logger.info(f"  - label: {label}")
    logger.info(f"  - is_unread: {is_unread}")
    logger.info(f"  - keywords: {keywords}")

    # Add base filters to exclude drafts and sent items, only include inbox
    query_parts.append("in:inbox -label:draft -label:sent")

    if sender:
        query_parts.append(f"from:{sender}")
    if recipient:
        query_parts.append(f"to:{recipient}")
    if subject:
        query_parts.append(f"subject:{subject}")
    if has_attachment:
        query_parts.append("has:attachment")
    if filename:
        query_parts.append(f"filename:{filename}")
    if after_date:
        date_str = after_date.strftime("%Y/%m/%d")
        query_parts.append(f"after:{date_str}")
    if before_date:
        date_str = before_date.strftime("%Y/%m/%d")
        query_parts.append(f"before:{date_str}")
    if newer_than_days:
        query_parts.append(f"newer_than:{newer_than_days}d")
    if older_than_days:
        query_parts.append(f"older_than:{older_than_days}d")
    if keywords:
        query_parts.append(f"({keywords})")
    if label:
        query_parts.append(f"label:{label}")
    if is_unread:
        query_parts.append("is:unread")

    query = " ".join(query_parts)
    logger.info(f"Search query: {query}")

    try:
        logger.info("Executing Gmail API search...")
        thread_contents = []
        next_page_token = None
        processed_threads = set()

        while True:
            request = service.users().messages().list(
                userId='me',
                q=query,
                pageToken=next_page_token,
                maxResults=100
            )
            response = request.execute()
            messages = response.get('messages', [])

            logger.info(f"Found {len(messages)} messages matching query")

            if not messages:
                logger.warning(
                    "No messages found matching the search criteria")
                return ""

            for message in messages:
                thread_id = message['threadId']

                if thread_id in processed_threads:
                    continue
                processed_threads.add(thread_id)

                thread = service.users().threads().get(
                    userId='me',
                    id=thread_id,
                    format='full'
                ).execute()

                thread_content = []

                # Process each message in the thread
                for msg in thread['messages']:
                    headers = msg['payload']['headers']
                    subject = next(
                        (h['value'] for h in headers if h['name'].lower() == 'subject'), 'No Subject')
                    from_header = next(
                        (h['value'] for h in headers if h['name'].lower() == 'from'), '')
                    date = next(
                        (h['value'] for h in headers if h['name'].lower() == 'date'), '')

                    # Extract message body
                    body = ""
                    if 'parts' in msg['payload']:
                        # Check each part for text/plain or text/html
                        for part in msg['payload']['parts']:
                            mime_type = part.get('mimeType')
                            if mime_type in ('text/plain', 'text/html'):
                                part_body = part.get(
                                    'body', {}).get('data', '')
                                if part_body:
                                    decoded_body = base64.urlsafe_b64decode(
                                        part_body).decode('utf-8', errors='replace')
                                    # Pick the part that has content
                                    if decoded_body.strip():
                                        body = decoded_body.strip()
                                        break
                    elif 'body' in msg['payload'] and 'data' in msg['payload']['body']:
                        # Fallback for single-part messages
                        body = base64.urlsafe_b64decode(
                            msg['payload']['body']['data']
                        ).decode('utf-8', errors='replace')

                    # Skip empty messages
                    if not body.strip():
                        continue

                    # Format message content
                    message_content = [
                        f"Date: {date}",
                        f"From: {from_header}",
                        f"Subject: {subject}",
                        "",
                        body.strip(),
                        ""  # Blank line between messages
                    ]
                    thread_content.extend(message_content)

                # Only add thread if it has content
                if thread_content:
                    thread_contents.append("\n".join(thread_content))

                if max_results and len(processed_threads) >= max_results:
                    return "\n\n".join(thread_contents)

            next_page_token = response.get('nextPageToken')
            if not next_page_token:
                break

        return "\n\n".join(thread_contents)

    except HttpError as error:
        logger.error(f"Gmail API error: {str(error)}")
        logger.error(
            f"Error details: {error.resp.status} - {error.resp.reason}")
        raise error


# Example usage:
"""
from datetime import datetime, timedelta

# Initialize your Gmail API service first
service = build("gmail", "v1", credentials=creds)

# Example searches:
# 1. Search for recent emails with attachments from a specific sender
results = search_emails(
    service,
    sender="example@gmail.com",
    has_attachment=True,
    newer_than_days=7
)

# 2. Search for unread emails with PDF attachments
results = search_emails(
    service,
    has_attachment=True,
    filename="pdf",
    is_unread=True
)

# 3. Search for emails within a date range
results = search_emails(
    service,
    after_date=datetime.now() - timedelta(days=30),
    before_date=datetime.now(),
    max_results=50
)

for email in results:
    print(f"Subject: {email['subject']}")
    print(f"From: {email['sender_name']} <{email['sender_email']}>")
    print(f"Date: {email['date']}")
    print(f"Attachments: {', '.join(email['attachments']) if email['attachments'] else 'None'}")
    print(f"Number of your responses: {email['response_count']}")
    print("-------------------")
"""
