import json
import os
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Optional
from googleapiclient.discovery import Resource
from pathlib import Path

logger = logging.getLogger(__name__)


class GmailIndexer:
    def __init__(self, service: Resource, index_file: str = "email_index.json"):
        """
        Initialize the Gmail indexer.

        Args:
            service: Authenticated Gmail API service instance
            index_file: Path to the index file (relative to the user's home directory)
        """
        self.service = service
        self.index_path = Path.home() / ".gmail_mcp" / index_file
        self.index_path.parent.mkdir(parents=True, exist_ok=True)

        # Get user's email address
        try:
            profile = self.service.users().getProfile(userId='me').execute()
            self.user_email = profile['emailAddress']
            logger.info(
                f"Initialized GmailIndexer for user: {self.user_email}")
        except Exception as e:
            logger.error(f"Error getting user email: {e}")
            self.user_email = None

        self.index: Dict[str, Dict] = self._load_index()

    def _load_index(self) -> Dict[str, Dict]:
        """Load the existing index or create a new one."""
        if self.index_path.exists():
            try:
                with open(self.index_path, 'r') as f:
                    return json.load(f)
            except json.JSONDecodeError:
                logger.warning("Corrupted index file, creating new index")
                return {"last_update": None, "messages": {}}
        return {"last_update": None, "messages": {}}

    def _save_index(self) -> None:
        """Save the current index to file."""
        with open(self.index_path, 'w') as f:
            json.dump(self.index, f, indent=2)

    def _is_domain_email(self, email_header: str) -> bool:
        """
        Check if an email is from our domain.
        Handles both plain email addresses and formatted headers (e.g., "Name <email@domain.com>")
        """
        try:
            # Extract domain from user's email
            user_domain = self.user_email.split('@')[1].lower()

            # Extract email from header if it contains a display name
            if '<' in email_header and '>' in email_header:
                email_address = email_header.split(
                    '<')[1].split('>')[0].lower()
            else:
                email_address = email_header.lower()

            # Check if the email address contains our domain
            return user_domain in email_address
        except Exception as e:
            logger.error(f"Error checking domain email: {e}")
            return False

    def _get_message_metadata(self, msg_id: str) -> Optional[Dict]:
        """Fetch metadata for a single message."""
        try:
            message = self.service.users().messages().get(
                userId='me',
                id=msg_id,
                format='metadata',
                metadataHeaders=['From', 'To', 'Subject', 'Date']
            ).execute()

            headers = {
                header['name']: header['value']
                for header in message['payload']['headers']
            }

            from_header = headers.get('From', '')

            # Extract email from header if it contains a display name
            if '<' in from_header and '>' in from_header:
                from_email = from_header.split('<')[1].split('>')[0].lower()
            else:
                from_email = from_header.lower()

            # Check if this is from the user's exact email
            is_from_user = self.user_email and (
                self.user_email.lower() == from_email or
                self.user_email.lower() in from_header.lower()
            )

            # Skip ALL messages from the user (no exceptions)
            if is_from_user:
                return None

            return {
                'id': message['id'],
                'threadId': message['threadId'],
                'from': from_header,
                'to': headers.get('To', ''),
                'subject': headers.get('Subject', ''),
                'date': headers.get('Date', ''),
                'snippet': message.get('snippet', ''),
                'labelIds': message.get('labelIds', []),
                'is_from_user': is_from_user
            }
        except Exception as e:
            logger.error(f"Error fetching message {msg_id}: {e}")
            return None

    def update_index(self) -> int:
        """
        Update the email index using expanding time windows.
        Returns number of new messages added.
        """
        current_time = datetime.now()
        window_hours = 1
        new_messages = 0
        last_update = self.index.get("last_update")

        while window_hours <= 720:  # Max 30 days back
            try:
                query_time = current_time - timedelta(hours=window_hours)
                query = ""  # Remove time-based filtering temporarily

                results = self.service.users().messages().list(
                    userId='me',
                    q=query,
                    maxResults=100
                ).execute()

                messages = results.get('messages', [])

                for message in messages:
                    msg_id = message['id']
                    if msg_id not in self.index['messages']:
                        metadata = self._get_message_metadata(msg_id)
                        if metadata:
                            self.index['messages'][msg_id] = metadata
                            new_messages += 1

                if 'nextPageToken' not in results:
                    window_hours *= 2

            except Exception as e:
                logger.error(f"Error updating index: {e}")
                break

        self.index["last_update"] = int(current_time.timestamp())
        self._save_index()

        if new_messages > 0:
            logger.info(
                f"Index update complete. Added {new_messages} new messages")

        return new_messages

    def search_index(self,
                     query: Optional[str] = None,
                     sender: Optional[str] = None,
                     recipient: Optional[str] = None,
                     subject: Optional[str] = None,
                     after_date: Optional[datetime] = None,
                     before_date: Optional[datetime] = None,
                     labels: Optional[List[str]] = None) -> List[Dict]:
        """
        Search the indexed messages using various criteria.

        Returns:
            List of matching message metadata
        """
        results = []

        for msg_id, metadata in self.index['messages'].items():
            # Apply filters
            if query and query.lower() not in metadata['snippet'].lower():
                continue
            if sender and sender.lower() not in metadata['from'].lower():
                continue
            if recipient and recipient.lower() not in metadata['to'].lower():
                continue
            if subject and subject.lower() not in metadata['subject'].lower():
                continue
            if labels and not any(label in metadata['labelIds'] for label in labels):
                continue

            # Parse and check dates
            try:
                msg_date = datetime.strptime(
                    metadata['date'], '%a, %d %b %Y %H:%M:%S %z')

                # Convert after_date and before_date to aware datetime if they're naive
                if after_date and after_date.tzinfo is None:
                    after_date = after_date.replace(tzinfo=msg_date.tzinfo)
                if before_date and before_date.tzinfo is None:
                    before_date = before_date.replace(tzinfo=msg_date.tzinfo)

                if after_date and msg_date < after_date:
                    continue
                if before_date and msg_date > before_date:
                    continue
            except ValueError:
                logger.warning(f"Could not parse date for message {msg_id}")
                continue

            results.append(metadata)

        return sorted(results, key=lambda x: x['date'], reverse=True)

    def get_new_emails(self) -> List[Dict]:
        """
        Returns any new emails found after updating the index.
        Excludes emails from the user.
        """
        new_count = self.update_index()
        if new_count > 0:
            # Get all messages that aren't from the user
            all_messages = [
                msg for msg in self.index['messages'].values()
                if not msg.get('is_from_user', False)
            ]

            # Sort by date, newest first
            sorted_messages = sorted(
                all_messages,
                key=lambda x: datetime.strptime(
                    x['date'], '%a, %d %b %Y %H:%M:%S %z'
                ),
                reverse=True
            )

            # Return the newest messages
            return sorted_messages[:new_count]
        return []


def get_or_update_index(service: Resource) -> GmailIndexer:
    """
    Helper function to get an indexer instance and update it.

    Returns:
        GmailIndexer instance with updated index
    """
    indexer = GmailIndexer(service)
    indexer.update_index()
    return indexer
