"""
Gmail tool for reading and managing Gmail emails.

Usage:
    from connectonion import Agent, Gmail

    gmail = Gmail()
    agent = Agent("assistant", tools=[gmail])

    # Agent can now use:
    # - read_inbox(last, unread)
    # - get_sent_emails(max_results)
    # - get_all_emails(max_results)
    # - search_emails(query, max_results)
    # - get_email_body(email_id)
    # - get_email_attachments(email_id)
    # - mark_read(email_id)
    # - mark_unread(email_id)
    # - archive_email(email_id)
    # - star_email(email_id)
    # - get_labels()
    # - add_label(email_id, label)
    # - get_emails_with_label(label, max_results)
    # - count_unread()

Example:
    from connectonion import Agent, Gmail

    gmail = Gmail()
    agent = Agent(
        name="gmail-assistant",
        system_prompt="You are a Gmail assistant.",
        tools=[gmail]
    )

    agent.input("Show me my recent emails")
    agent.input("Search for emails from alice@example.com")
"""

import os
import base64
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError


class Gmail:
    """Gmail tool for reading and managing emails."""

    def __init__(self):
        """Initialize Gmail tool.

        Validates that gmail.readonly scope is authorized.
        Raises ValueError if scope is missing.
        """
        scopes = os.getenv("GOOGLE_SCOPES", "")
        if "gmail.readonly" not in scopes:
            raise ValueError(
                "Missing 'gmail.readonly' scope.\n"
                f"Current scopes: {scopes}\n"
                "Please authorize Gmail read access:\n"
                "  co auth google"
            )

        self._service = None

    def _get_service(self):
        """Get Gmail API service (lazy load)."""
        if self._service:
            return self._service

        access_token = os.getenv("GOOGLE_ACCESS_TOKEN")
        refresh_token = os.getenv("GOOGLE_REFRESH_TOKEN")
        expires_at_str = os.getenv("GOOGLE_TOKEN_EXPIRES_AT")

        if not access_token or not refresh_token:
            raise ValueError(
                "Google OAuth credentials not found.\n"
                "Run: co auth google"
            )

        # Check if token is expired or about to expire (within 5 minutes)
        if expires_at_str:
            from datetime import datetime, timedelta
            expires_at = datetime.fromisoformat(expires_at_str.replace('Z', '+00:00'))
            now = datetime.utcnow().replace(tzinfo=expires_at.tzinfo) if expires_at.tzinfo else datetime.utcnow()

            if now >= expires_at - timedelta(minutes=5):
                # Token expired or about to expire, refresh via backend
                access_token = self._refresh_via_backend(refresh_token)

        # Create credentials without client_id/client_secret
        # Backend handles token refresh, so we don't need auto-refresh
        creds = Credentials(
            token=access_token,
            refresh_token=refresh_token,
            token_uri="https://oauth2.googleapis.com/token",
            client_id=None,
            client_secret=None,
            scopes=["https://www.googleapis.com/auth/gmail.readonly",
                   "https://www.googleapis.com/auth/gmail.modify"]
        )

        self._service = build('gmail', 'v1', credentials=creds)
        return self._service

    def _refresh_via_backend(self, refresh_token: str) -> str:
        """Refresh access token via backend API.

        Args:
            refresh_token: The refresh token

        Returns:
            New access token
        """
        import httpx

        # Get backend URL and auth
        backend_url = os.getenv("OPENONION_API_URL", "https://o.openonion.ai")
        api_key = os.getenv("OPENONION_API_KEY")

        if not api_key:
            raise ValueError(
                "OPENONION_API_KEY not found.\n"
                "This is needed to refresh tokens via backend."
            )

        # Call backend refresh endpoint
        response = httpx.post(
            f"{backend_url}/api/v1/oauth/google/refresh",
            headers={"Authorization": f"Bearer {api_key}"},
            json={"refresh_token": refresh_token}
        )

        if response.status_code != 200:
            raise ValueError(
                f"Failed to refresh token via backend: {response.text}"
            )

        data = response.json()
        new_access_token = data["access_token"]
        expires_at = data["expires_at"]

        # Update environment variables for this session
        os.environ["GOOGLE_ACCESS_TOKEN"] = new_access_token
        os.environ["GOOGLE_TOKEN_EXPIRES_AT"] = expires_at

        # Update .env file if it exists
        env_file = os.path.join(os.getenv("AGENT_CONFIG_PATH", os.path.expanduser("~/.co")), "keys.env")
        if os.path.exists(env_file):
            with open(env_file, 'r') as f:
                lines = f.readlines()

            with open(env_file, 'w') as f:
                for line in lines:
                    if line.startswith("GOOGLE_ACCESS_TOKEN="):
                        f.write(f"GOOGLE_ACCESS_TOKEN={new_access_token}\n")
                    elif line.startswith("GOOGLE_TOKEN_EXPIRES_AT="):
                        f.write(f"GOOGLE_TOKEN_EXPIRES_AT={expires_at}\n")
                    else:
                        f.write(line)

        return new_access_token

    def _format_emails(self, messages, max_results=10):
        """Helper to format email list."""
        if not messages:
            return "No emails found."

        service = self._get_service()
        emails = []

        for msg in messages[:max_results]:
            message = service.users().messages().get(
                userId='me',
                id=msg['id'],
                format='metadata',
                metadataHeaders=['From', 'Subject', 'Date']
            ).execute()

            headers = message['payload']['headers']
            subject = next((h['value'] for h in headers if h['name'] == 'Subject'), 'No Subject')
            from_email = next((h['value'] for h in headers if h['name'] == 'From'), 'Unknown')
            date = next((h['value'] for h in headers if h['name'] == 'Date'), 'Unknown')

            snippet = message.get('snippet', '')
            is_unread = 'UNREAD' in message.get('labelIds', [])

            emails.append({
                'id': msg['id'],
                'from': from_email,
                'subject': subject,
                'date': date,
                'snippet': snippet,
                'unread': is_unread
            })

        # Format output
        output = [f"Found {len(emails)} email(s):\n"]
        for i, email in enumerate(emails, 1):
            status = "[UNREAD]" if email['unread'] else ""
            output.append(f"{i}. {status} From: {email['from']}")
            output.append(f"   Subject: {email['subject']}")
            output.append(f"   Date: {email['date']}")
            output.append(f"   Preview: {email['snippet'][:80]}...")
            output.append(f"   ID: {email['id']}\n")

        return "\n".join(output)

    # === Reading ===

    def read_inbox(self, last: int = 10, unread: bool = False) -> str:
        """Read emails from inbox.

        Args:
            last: Number of emails to retrieve (default: 10)
            unread: Only get unread emails (default: False)

        Returns:
            Formatted string with email list
        """
        service = self._get_service()

        query = "is:unread in:inbox" if unread else "in:inbox"

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

        messages = results.get('messages', [])
        return self._format_emails(messages, last)

    def get_sent_emails(self, max_results: int = 10) -> str:
        """Get emails you sent.

        Args:
            max_results: Number of emails to retrieve (default: 10)

        Returns:
            Formatted string with sent email list
        """
        service = self._get_service()

        results = service.users().messages().list(
            userId='me',
            q="in:sent",
            maxResults=max_results
        ).execute()

        messages = results.get('messages', [])
        return self._format_emails(messages, max_results)

    def get_all_emails(self, max_results: int = 50) -> str:
        """Get emails from all folders.

        Args:
            max_results: Number of emails to retrieve (default: 50)

        Returns:
            Formatted string with email list
        """
        service = self._get_service()

        results = service.users().messages().list(
            userId='me',
            maxResults=max_results
        ).execute()

        messages = results.get('messages', [])
        return self._format_emails(messages, max_results)

    # === Search ===

    def search_emails(self, query: str, max_results: int = 10) -> str:
        """Search emails using Gmail query syntax.

        Args:
            query: Gmail search query (e.g., "from:alice@example.com", "subject:meeting")
            max_results: Number of results to return (default: 10)

        Returns:
            Formatted string with matching emails
        """
        service = self._get_service()

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

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

        if not messages:
            return f"No emails found matching query: {query}"

        return self._format_emails(messages, max_results)

    # === Content ===

    def get_email_body(self, email_id: str) -> str:
        """Get full email body.

        Args:
            email_id: Gmail message ID

        Returns:
            Full email content with headers
        """
        service = self._get_service()

        message = service.users().messages().get(
            userId='me',
            id=email_id,
            format='full'
        ).execute()

        headers = message['payload']['headers']
        subject = next((h['value'] for h in headers if h['name'] == 'Subject'), 'No Subject')
        from_email = next((h['value'] for h in headers if h['name'] == 'From'), 'Unknown')
        to_email = next((h['value'] for h in headers if h['name'] == 'To'), 'Unknown')
        date = next((h['value'] for h in headers if h['name'] == 'Date'), 'Unknown')

        # Extract body
        body = ""
        if 'parts' in message['payload']:
            for part in message['payload']['parts']:
                if part['mimeType'] == 'text/plain':
                    if 'data' in part['body']:
                        body = base64.urlsafe_b64decode(part['body']['data']).decode('utf-8')
                        break
        elif 'body' in message['payload'] and 'data' in message['payload']['body']:
            body = base64.urlsafe_b64decode(message['payload']['body']['data']).decode('utf-8')

        if not body:
            body = message.get('snippet', 'No body content')

        output = [
            f"From: {from_email}",
            f"To: {to_email}",
            f"Subject: {subject}",
            f"Date: {date}",
            "\n--- Email Body ---\n",
            body
        ]

        return "\n".join(output)

    def get_email_attachments(self, email_id: str) -> str:
        """List attachments in email.

        Args:
            email_id: Gmail message ID

        Returns:
            List of attachment names and sizes
        """
        service = self._get_service()

        message = service.users().messages().get(
            userId='me',
            id=email_id,
            format='full'
        ).execute()

        attachments = []

        if 'parts' in message['payload']:
            for part in message['payload']['parts']:
                if part.get('filename'):
                    size = part['body'].get('size', 0)
                    attachments.append({
                        'filename': part['filename'],
                        'size': size,
                        'id': part['body'].get('attachmentId', '')
                    })

        if not attachments:
            return "No attachments in this email."

        output = [f"Found {len(attachments)} attachment(s):\n"]
        for i, att in enumerate(attachments, 1):
            size_kb = att['size'] / 1024
            output.append(f"{i}. {att['filename']} ({size_kb:.1f} KB)")
            output.append(f"   ID: {att['id']}\n")

        return "\n".join(output)

    # === Actions ===

    def mark_read(self, email_id: str) -> str:
        """Mark email as read.

        Args:
            email_id: Gmail message ID

        Returns:
            Confirmation message
        """
        service = self._get_service()

        service.users().messages().modify(
            userId='me',
            id=email_id,
            body={'removeLabelIds': ['UNREAD']}
        ).execute()

        return f"Marked email as read: {email_id}"

    def mark_unread(self, email_id: str) -> str:
        """Mark email as unread.

        Args:
            email_id: Gmail message ID

        Returns:
            Confirmation message
        """
        service = self._get_service()

        service.users().messages().modify(
            userId='me',
            id=email_id,
            body={'addLabelIds': ['UNREAD']}
        ).execute()

        return f"Marked email as unread: {email_id}"

    def archive_email(self, email_id: str) -> str:
        """Archive email (remove from inbox).

        Args:
            email_id: Gmail message ID

        Returns:
            Confirmation message
        """
        service = self._get_service()

        service.users().messages().modify(
            userId='me',
            id=email_id,
            body={'removeLabelIds': ['INBOX']}
        ).execute()

        return f"Archived email: {email_id}"

    def star_email(self, email_id: str) -> str:
        """Add star to email.

        Args:
            email_id: Gmail message ID

        Returns:
            Confirmation message
        """
        service = self._get_service()

        service.users().messages().modify(
            userId='me',
            id=email_id,
            body={'addLabelIds': ['STARRED']}
        ).execute()

        return f"Starred email: {email_id}"

    # === Labels ===

    def get_labels(self) -> str:
        """List all Gmail labels.

        Returns:
            List of label names and IDs
        """
        service = self._get_service()

        results = service.users().labels().list(userId='me').execute()
        labels = results.get('labels', [])

        if not labels:
            return "No labels found."

        output = [f"Found {len(labels)} label(s):\n"]
        for label in labels:
            label_type = label.get('type', 'user')
            output.append(f"- {label['name']} (ID: {label['id']}, Type: {label_type})")

        return "\n".join(output)

    def add_label(self, email_id: str, label: str) -> str:
        """Add label to email.

        Args:
            email_id: Gmail message ID
            label: Label name or ID

        Returns:
            Confirmation message
        """
        service = self._get_service()

        # Try to find label by name first
        results = service.users().labels().list(userId='me').execute()
        labels = results.get('labels', [])

        label_id = label
        for lbl in labels:
            if lbl['name'].lower() == label.lower():
                label_id = lbl['id']
                break

        service.users().messages().modify(
            userId='me',
            id=email_id,
            body={'addLabelIds': [label_id]}
        ).execute()

        return f"Added label '{label}' to email: {email_id}"

    def get_emails_with_label(self, label: str, max_results: int = 10) -> str:
        """Get emails with specific label.

        Args:
            label: Label name (e.g., "Important", "Work")
            max_results: Number of emails to retrieve (default: 10)

        Returns:
            Formatted string with email list
        """
        service = self._get_service()

        # Find label ID by name
        results = service.users().labels().list(userId='me').execute()
        labels = results.get('labels', [])

        label_id = None
        for lbl in labels:
            if lbl['name'].lower() == label.lower():
                label_id = lbl['id']
                break

        if not label_id:
            return f"Label not found: {label}"

        results = service.users().messages().list(
            userId='me',
            labelIds=[label_id],
            maxResults=max_results
        ).execute()

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

        if not messages:
            return f"No emails with label: {label}"

        return self._format_emails(messages, max_results)

    # === Stats ===

    def count_unread(self) -> str:
        """Count unread emails.

        Returns:
            Number of unread emails
        """
        service = self._get_service()

        results = service.users().messages().list(
            userId='me',
            q="is:unread",
            maxResults=1
        ).execute()

        # Get total from resultSizeEstimate
        count = results.get('resultSizeEstimate', 0)

        return f"You have {count} unread email(s)."
