"""Gmail connector via OAuth2."""

import base64
import logging
import os

from dotenv import load_dotenv
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build

logger = logging.getLogger(__name__)

TOKEN_URI = "https://oauth2.googleapis.com/token"


class GmailConnector:
    """Gmail connector authenticated via OAuth2 refresh token.

    Loads credentials from constructor args or falls back to environment
    variables (GMAIL_CLIENT_ID, GMAIL_CLIENT_SECRET, GMAIL_REFRESH_TOKEN).
    """

    def __init__(
        self,
        client_id: str | None = None,
        client_secret: str | None = None,
        refresh_token: str | None = None,
    ) -> None:
        """Initialize the connector and build the Gmail API service.

        Args:
            client_id: OAuth2 client ID. Falls back to GMAIL_CLIENT_ID env var.
            client_secret: OAuth2 client secret. Falls back to GMAIL_CLIENT_SECRET env var.
            refresh_token: OAuth2 refresh token. Falls back to GMAIL_REFRESH_TOKEN env var.
        """
        load_dotenv()
        client_id = client_id or os.environ["GMAIL_CLIENT_ID"]
        client_secret = client_secret or os.environ["GMAIL_CLIENT_SECRET"]
        refresh_token = refresh_token or os.environ["GMAIL_REFRESH_TOKEN"]

        creds = Credentials(
            token=None,
            refresh_token=refresh_token,
            client_id=client_id,
            client_secret=client_secret,
            token_uri=TOKEN_URI,
        )
        self._service = build("gmail", "v1", credentials=creds)
        logger.info("GmailConnector initialized")

    def list_labels(self) -> list[dict]:
        """List all labels (folders) in the mailbox.

        Returns:
            A list of label dicts with keys like 'id', 'name', 'type'.
        """
        resp = self._service.users().labels().list(userId="me").execute()
        return resp.get("labels", [])

    def list_messages(
        self,
        label_ids: list[str] | None = None,
        query: str | None = None,
        max_results: int = 10,
    ) -> list[dict]:
        """List message stubs matching the given filters.

        Args:
            label_ids: Filter by label IDs (e.g. ["INBOX"]).
            query: Gmail search query (e.g. "from:alice@test.com").
            max_results: Maximum number of messages to return.

        Returns:
            A list of message stub dicts with 'id' and 'threadId'.
        """
        kwargs: dict = {"userId": "me", "maxResults": max_results}
        if label_ids:
            kwargs["labelIds"] = label_ids
        if query:
            kwargs["q"] = query
        resp = self._service.users().messages().list(**kwargs).execute()
        return resp.get("messages", [])

    def get_message(self, message_id: str) -> dict:
        """Fetch a full message and decode its body.

        Args:
            message_id: The Gmail message ID.

        Returns:
            The full message dict with an added 'body_text' key
            containing the decoded plain-text body.
        """
        msg = (
            self._service.users()
            .messages()
            .get(userId="me", id=message_id, format="full")
            .execute()
        )
        msg["body_text"] = _decode_body(msg.get("payload", {}))
        return msg

    def modify_message(
        self,
        message_id: str,
        add_label_ids: list[str] | None = None,
        remove_label_ids: list[str] | None = None,
    ) -> dict:
        """Modify a message's labels.

        Args:
            message_id: The Gmail message ID.
            add_label_ids: Label IDs to add (e.g. ["STARRED"]).
            remove_label_ids: Label IDs to remove (e.g. ["UNREAD"]).

        Returns:
            The updated message dict.
        """
        body: dict = {}
        if add_label_ids:
            body["addLabelIds"] = add_label_ids
        if remove_label_ids:
            body["removeLabelIds"] = remove_label_ids
        return self._service.users().messages().modify(userId="me", id=message_id, body=body).execute()

    def mark_as_read(self, message_id: str) -> dict:
        """Mark a message as read."""
        return self.modify_message(message_id, remove_label_ids=["UNREAD"])

    def mark_as_unread(self, message_id: str) -> dict:
        """Mark a message as unread."""
        return self.modify_message(message_id, add_label_ids=["UNREAD"])

    def archive(self, message_id: str) -> dict:
        """Archive a message (remove from INBOX)."""
        return self.modify_message(message_id, remove_label_ids=["INBOX"])

    def trash(self, message_id: str) -> dict:
        """Move a message to trash."""
        return self._service.users().messages().trash(userId="me", id=message_id).execute()

    def untrash(self, message_id: str) -> dict:
        """Remove a message from trash."""
        return self._service.users().messages().untrash(userId="me", id=message_id).execute()

    def list_attachments(self, message_id: str) -> list[dict]:
        """List attachments for a message.

        Args:
            message_id: The Gmail message ID.

        Returns:
            A list of dicts with keys: filename, mimeType, size, attachmentId.
        """
        msg = self.get_message(message_id)
        return _collect_attachments(msg.get("payload", {}))

    def download_attachment(self, message_id: str, attachment_id: str) -> bytes:
        """Download a single attachment by ID.

        Args:
            message_id: The Gmail message ID.
            attachment_id: The attachment ID from list_attachments().

        Returns:
            The raw attachment bytes.
        """
        resp = (
            self._service.users()
            .messages()
            .attachments()
            .get(userId="me", messageId=message_id, id=attachment_id)
            .execute()
        )
        return base64.urlsafe_b64decode(resp["data"])

    def send_message(self, to: str, subject: str, body: str) -> dict:
        """Send a plain-text email.

        Args:
            to: Recipient email address.
            subject: Email subject line.
            body: Plain-text body content.

        Returns:
            The sent message dict with 'id' and 'threadId'.
        """
        from email.message import EmailMessage

        msg = EmailMessage()
        msg["To"] = to
        msg["Subject"] = subject
        msg.set_content(body)

        raw = base64.urlsafe_b64encode(msg.as_bytes()).decode()
        return self._service.users().messages().send(userId="me", body={"raw": raw}).execute()

    def __repr__(self) -> str:
        return "GmailConnector()"


def _collect_attachments(payload: dict) -> list[dict]:
    """Walk the MIME tree and collect attachment metadata."""
    attachments = []
    filename = payload.get("filename")
    if filename:
        body = payload.get("body", {})
        attachments.append(
            {
                "filename": filename,
                "mimeType": payload.get("mimeType", ""),
                "size": body.get("size", 0),
                "attachmentId": body.get("attachmentId", ""),
            }
        )
    for part in payload.get("parts", []):
        attachments.extend(_collect_attachments(part))
    return attachments


def _decode_body(payload: dict) -> str:
    """Extract and decode the plain-text body from a message payload.

    Walks the MIME tree recursively to find the first text/plain part.

    Args:
        payload: The message payload dict from Gmail API.

    Returns:
        The decoded plain-text body, or empty string if none found.
    """
    mime_type = payload.get("mimeType", "")
    body_data = payload.get("body", {}).get("data")

    if mime_type == "text/plain" and body_data:
        return base64.urlsafe_b64decode(body_data).decode()

    for part in payload.get("parts", []):
        text = _decode_body(part)
        if text:
            return text

    return ""
