import os
import json
import base64
from datetime import datetime
from typing import Dict, List
from googleapiclient.errors import HttpError
from dotenv import load_dotenv
import email.utils
import time
from tqdm import tqdm
import logging

# Load environment variables
load_dotenv()

# Get environment variables
USER_EMAIL = os.getenv('USER_EMAIL')
OUTPUT_DIR = os.getenv('EMAIL_THREADS_DIR', 'email_threads')
MAX_RESULTS = int(os.getenv('MAX_RESULTS', 500))

logger = logging.getLogger("gmail_mcp")


def parse_email_date(date_str: str) -> datetime:
    """Parse email date string to datetime object."""
    try:
        timestamp = email.utils.mktime_tz(email.utils.parsedate_tz(date_str))
        return datetime.fromtimestamp(timestamp)
    except Exception:
        return datetime.now()


def process_thread(service, message_id: str, output_dir: str) -> None:
    """Process a single email thread."""
    try:
        # Get full message details
        msg = service.users().messages().get(
            userId='me',
            id=message_id,
            format='full'
        ).execute()

        thread_id = msg['threadId']

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

        # Process the conversation thread
        conversation = []
        participants = set()
        subject = 'No Subject'  # Default subject

        for thread_message in thread['messages']:
            headers = thread_message['payload']['headers']
            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'), '')
            subject = next(
                (h['value'] for h in headers if h['name'].lower() == 'subject'), subject)

            # Get message body
            body = '[No body content]'
            if 'parts' in thread_message['payload']:
                for part in thread_message['payload']['parts']:
                    if part['mimeType'] == 'text/plain':
                        if 'data' in part['body']:
                            body = base64.urlsafe_b64decode(
                                part['body']['data']
                            ).decode('utf-8', errors='replace')
                            break
            elif 'body' in thread_message['payload'] and 'data' in thread_message['payload']['body']:
                body = base64.urlsafe_b64decode(
                    thread_message['payload']['body']['data']
                ).decode('utf-8', errors='replace')

            # Determine if this is from me or the other person
            is_from_me = USER_EMAIL in from_header.lower()

            # Extract email addresses
            if '<' in from_header:
                email_addr = from_header.split('<')[1].rstrip('>')
                name = from_header.split('<')[0].strip()
            else:
                email_addr = from_header
                name = email_addr

            participants.add(email_addr)

            # Add to conversation
            conversation.append({
                'from': 'me' if is_from_me else name,
                'email': email_addr,
                'date': date,
                'content': body,
                'message_id': thread_message['id']
            })

        # Sort conversation by date
        conversation.sort(key=lambda x: parse_email_date(x['date']))

        # Create thread data
        thread_data = {
            'subject': subject,
            'participants': list(participants),
            'messages': conversation,
            'last_updated': datetime.now().isoformat()
        }

        # Save thread file
        thread_filename = f"thread_{thread_id}.json"
        thread_path = os.path.join(output_dir, thread_filename)

        with open(thread_path, 'w', encoding='utf-8') as f:
            json.dump(thread_data, f, indent=2, ensure_ascii=False)

    except Exception as e:
        logger.error(f"Error processing message {message_id}: {str(e)}")


def setup_email_index(service, output_dir: str = OUTPUT_DIR) -> None:
    """Index all sent messages and create conversation files organized by thread."""
    try:
        start_time = time.time()

        # Create output directory if it doesn't exist
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)

        logger.info("Fetching sent messages...")
        results = service.users().messages().list(
            userId='me',
            labelIds=['SENT'],
            maxResults=MAX_RESULTS
        ).execute()

        messages = results.get('messages', [])
        total_messages = len(messages)
        logger.info(f"Found {total_messages} messages to process...")

        # Process each message sequentially with progress bar
        for message in tqdm(messages, desc="Processing threads"):
            process_thread(service, message['id'], output_dir)

        end_time = time.time()
        duration = end_time - start_time

        logger.info(f"\nIndexing complete!")
        logger.info(
            f"Processed {total_messages} messages in {duration:.2f} seconds")
        logger.info(f"Files saved in: {os.path.abspath(output_dir)}")

    except HttpError as error:
        logger.error(f"An error occurred: {error}")


if __name__ == "__main__":
    from gmail_mcp.utils.GCP.gmail_auth import get_gmail_service

    service = get_gmail_service()
    setup_email_index(service)
