"""Gmail mailbox retrieval handler via OAuth2.

Pulls mail from a shared intake mailbox in Gmail. Each tenant only sees
messages that carry that tenant's label. Emails are returned as complete
raw RFC822 (.eml) files, base64-encoded.

Auth model:
    OAuth2 client_id/client_secret + refresh_token.
    The refresh_token was granted by the mailbox owner.
    We exchange refresh_token -> access_token on demand.

Isolation model:
    Settings must provide a TENANT_NAME which maps to a Gmail label
    (format: TENANT/{TENANT_NAME}). Only messages carrying that label
    are listed/read.

Usage::

    from entr_adapter_core.connectors.gmail import (
        EntrMailboxRetrievalHandler,
        GmailConnectorConfig,
    )

    config = GmailConnectorConfig(
        tenant_name="Acme",
        gmail_client_id="...",
        gmail_client_secret="...",
        gmail_refresh_token="...",
    )
    handler = EntrMailboxRetrievalHandler(config)

Configuration can also be provided via any object that exposes the
required attributes (TENANT_NAME, GMAIL_CLIENT_ID, GMAIL_CLIENT_SECRET,
GMAIL_REFRESH_TOKEN) — e.g. a pydantic-settings model or a plain
namespace. The ``GmailConnectorConfig`` Pydantic model is provided for
convenience and validation.
"""

import base64
import logging
from datetime import date
from typing import List, Optional
from urllib.parse import quote

try:
    import httpx
except ImportError as exc:
    raise ImportError(
        "The Gmail connector requires the 'httpx' package. "
        "Install it with: uv add entr-adapter-core[gmail]"
    ) from exc

from pydantic import BaseModel, Field

from entr_adapter_core.handlers import RetrievalHandler
from entr_adapter_core.schemas import (
    DocumentDownloadResponse,
    DocumentIdentifier,
    RetrievalPayload,
    RetrievalResponse,
    StatusEnum,
)

logger = logging.getLogger(__name__)


class GmailConnectorConfig(BaseModel):
    """Configuration for the Gmail mailbox retrieval connector.

    Tenant-specific settings that must be provided at instantiation.
    All fields map to OAuth2 credentials and Gmail filtering options.

    Attributes:
        tenant_name: Identifier for the tenant. Used to derive the Gmail
            label (``TENANT/{tenant_name}``) unless ``label_override`` is set.
        gmail_client_id: OAuth2 client ID from Google Cloud Console.
        gmail_client_secret: OAuth2 client secret from Google Cloud Console.
        gmail_refresh_token: Long-lived refresh token granted by the
            mailbox owner during the OAuth2 consent flow.
        label_override: Optional custom Gmail label name. When set, the
            connector uses this label instead of ``TENANT/{tenant_name}``.
        max_results_per_page: Maximum messages per Gmail API page request.
            Gmail caps this at 500.
    """

    tenant_name: str = Field(
        ...,
        description="Tenant identifier, used to derive the Gmail label (TENANT/{tenant_name})",
    )
    gmail_client_id: str = Field(
        ...,
        description="OAuth2 client ID from Google Cloud Console",
    )
    gmail_client_secret: str = Field(
        ...,
        description="OAuth2 client secret from Google Cloud Console",
    )
    gmail_refresh_token: str = Field(
        ...,
        description="Long-lived OAuth2 refresh token granted by the mailbox owner",
    )
    label_override: Optional[str] = Field(
        default=None,
        description="Custom Gmail label name; overrides the default TENANT/{tenant_name}",
    )
    max_results_per_page: int = Field(
        default=500,
        ge=1,
        le=500,
        description="Maximum messages per Gmail API page request (max 500)",
    )

    # Expose uppercase aliases so the handler can read either style.
    @property
    def TENANT_NAME(self) -> str:  # noqa: N802
        return self.tenant_name

    @property
    def GMAIL_CLIENT_ID(self) -> str:  # noqa: N802
        return self.gmail_client_id

    @property
    def GMAIL_CLIENT_SECRET(self) -> str:  # noqa: N802
        return self.gmail_client_secret

    @property
    def GMAIL_REFRESH_TOKEN(self) -> str:  # noqa: N802
        return self.gmail_refresh_token


class EntrMailboxRetrievalHandler(RetrievalHandler):
    """RetrievalHandler that pulls mail from a Gmail inbox.

    Accepts either a ``GmailConnectorConfig`` instance or any settings object
    that exposes ``TENANT_NAME``, ``GMAIL_CLIENT_ID``, ``GMAIL_CLIENT_SECRET``,
    and ``GMAIL_REFRESH_TOKEN`` attributes.

    Args:
        settings: Configuration object with tenant-specific Gmail credentials.
            Can be a ``GmailConnectorConfig``, a pydantic-settings model, or
            any object with the required attributes.
        label_override: Use a custom Gmail label instead of the default
            ``TENANT/{TENANT_NAME}``. Ignored when ``settings`` is a
            ``GmailConnectorConfig`` with its own ``label_override``.

    Example::

        config = GmailConnectorConfig(
            tenant_name="Acme",
            gmail_client_id="123.apps.googleusercontent.com",
            gmail_client_secret="secret",
            gmail_refresh_token="1//refresh",
        )
        handler = EntrMailboxRetrievalHandler(config)
        response = await handler.list_documents(payload)
    """

    GMAIL_TOKEN_URL = "https://oauth2.googleapis.com/token"
    GMAIL_API_BASE = "https://gmail.googleapis.com/gmail/v1"

    def __init__(self, settings, label_override: str = None):
        super().__init__(settings)
        # If settings is a GmailConnectorConfig, prefer its label_override
        if isinstance(settings, GmailConnectorConfig) and settings.label_override:
            self._label_override = settings.label_override
        else:
            self._label_override = label_override
        logger.info(
            "Initialized EntrMailboxRetrievalHandler for tenant: %s%s",
            self.settings.TENANT_NAME,
            f" (label override: {self._label_override})" if self._label_override else "",
        )
        self._label_id_cache: Optional[str] = None

    # ── Internal helpers ─────────────────────────────────────────────────

    def _build_label_name(self) -> str:
        """Map tenant name -> Gmail label name (e.g. 'TENANT/Zuidberg')."""
        if self._label_override:
            return self._label_override
        return f"TENANT/{self.settings.TENANT_NAME}"

    async def _get_access_token(self, client: httpx.AsyncClient) -> str:
        """Exchange stored refresh token for a fresh access token."""
        data = {
            "client_id": self.settings.GMAIL_CLIENT_ID,
            "client_secret": self.settings.GMAIL_CLIENT_SECRET,
            "refresh_token": self.settings.GMAIL_REFRESH_TOKEN,
            "grant_type": "refresh_token",
        }
        resp = await client.post(self.GMAIL_TOKEN_URL, data=data, timeout=10)
        if resp.status_code != 200:
            logger.error("Failed to refresh access token: %d %s", resp.status_code, resp.text)
            raise RuntimeError("OAuth token refresh failed")
        return resp.json()["access_token"]

    def _auth_headers(self, access_token: str) -> dict:
        return {
            "Authorization": f"Bearer {access_token}",
            "Accept": "application/json",
        }

    async def _get_label_id(self, client: httpx.AsyncClient, access_token: str) -> str:
        """Resolve Gmail internal label ID for this tenant's label name."""
        if self._label_id_cache:
            return self._label_id_cache

        desired_label = self._build_label_name()
        url = f"{self.GMAIL_API_BASE}/users/me/labels"
        headers = self._auth_headers(access_token)
        resp = await client.get(url, headers=headers, timeout=10)

        if resp.status_code != 200:
            logger.error("Failed to list labels: %d %s", resp.status_code, resp.text)
            raise RuntimeError("Could not list Gmail labels")

        labels = resp.json().get("labels", [])
        for lbl in labels:
            if lbl.get("name") == desired_label:
                self._label_id_cache = lbl.get("id")
                return self._label_id_cache

        logger.error("Gmail label '%s' not found", desired_label)
        raise RuntimeError(f"Gmail label '{desired_label}' not found")

    async def _list_messages_for_label(
        self,
        client: httpx.AsyncClient,
        access_token: str,
        label_id: str,
        date_from: Optional[date] = None,
    ) -> List[dict]:
        """Return all message stubs for mail with this label, with pagination."""
        max_results = getattr(self.settings, "max_results_per_page", 500)
        base_url = f"{self.GMAIL_API_BASE}/users/me/messages?labelIds={label_id}&maxResults={max_results}"

        if date_from:
            date_str = date_from.strftime("%Y/%m/%d")
            query = f"after:{date_str}"
            base_url += f"&q={quote(query)}"
            logger.info("Gmail query filter: %s", query)

        headers = self._auth_headers(access_token)
        all_messages = []
        next_page_token = None

        while True:
            url = base_url
            if next_page_token:
                url += f"&pageToken={next_page_token}"

            resp = await client.get(url, headers=headers, timeout=10)
            if resp.status_code != 200:
                logger.error("Failed to list messages: %d %s", resp.status_code, resp.text)
                raise RuntimeError("Could not list Gmail messages for label")

            data = resp.json()
            messages = data.get("messages", []) or []
            all_messages.extend(messages)

            next_page_token = data.get("nextPageToken")
            if not next_page_token:
                break

            logger.info("Fetching next page, %d messages so far...", len(all_messages))

        return all_messages

    async def _get_message_raw_rfc822(
        self, client: httpx.AsyncClient, access_token: str, message_id: str
    ) -> bytes:
        """Fetch the full raw RFC822 email from Gmail (format=raw)."""
        url = f"{self.GMAIL_API_BASE}/users/me/messages/{message_id}?format=raw"
        headers = self._auth_headers(access_token)
        resp = await client.get(url, headers=headers, timeout=10)

        if resp.status_code != 200:
            logger.error("Failed to get raw message %s: %d %s", message_id, resp.status_code, resp.text)
            raise RuntimeError("Could not get raw Gmail message")

        raw_b64url = resp.json().get("raw")
        if not raw_b64url:
            raise RuntimeError("No raw content in Gmail response")

        # Gmail returns base64url (- and _ instead of + and /, no padding).
        b64std = raw_b64url.replace("-", "+").replace("_", "/")
        padding_needed = (4 - (len(b64std) % 4)) % 4
        b64std += "=" * padding_needed
        return base64.b64decode(b64std)

    # ── Public API ───────────────────────────────────────────────────────

    async def list_documents(self, payload: RetrievalPayload) -> RetrievalResponse:
        """Return recent emails for this tenant's label.

        Each email becomes a DocumentIdentifier with:
        - document_identifier = Gmail message ID
        - filename = Message ID with .eml extension
        """
        logger.info(
            "EntrMailboxRetrievalHandler.list_documents "
            "tenant=%s source=%s date_from=%s",
            self.settings.TENANT_NAME, payload.source_identifier, payload.date_from,
        )

        async with httpx.AsyncClient() as client:
            access_token = await self._get_access_token(client)
            label_id = await self._get_label_id(client, access_token)

            msgs = await self._list_messages_for_label(
                client, access_token, label_id,
                date_from=payload.date_from,
            )

            logger.info("Gmail returned %d emails matching filters", len(msgs))

            docs: List[DocumentIdentifier] = [
                DocumentIdentifier(
                    document_identifier=stub["id"],
                    filename=f"{stub['id']}.eml",
                )
                for stub in msgs
            ]

        return RetrievalResponse(
            status=StatusEnum.SUCCESS,
            message=f"Found {len(docs)} emails for tenant '{self._build_label_name()}'.",
            documents=docs,
        )

    async def get_document(self, document_identifier: str) -> DocumentDownloadResponse:
        """Return the full .eml of a Gmail message, base64-encoded."""
        logger.info(
            "EntrMailboxRetrievalHandler.get_document "
            "tenant=%s message_id=%s",
            self.settings.TENANT_NAME, document_identifier,
        )

        async with httpx.AsyncClient() as client:
            access_token = await self._get_access_token(client)
            raw_bytes = await self._get_message_raw_rfc822(client, access_token, document_identifier)

        eml_b64 = base64.b64encode(raw_bytes).decode("utf-8")
        filename = f"{document_identifier}.eml"

        return DocumentDownloadResponse(
            document_identifier=filename,
            adapter_document_id=document_identifier,
            content_b64=eml_b64,
        )
