import base64
from datetime import date
from typing import List, Optional
import httpx
from app.retrieval_engine.base import RetrievalHandler
from app.schemas import (
    RetrievalPayload,
    RetrievalResponse,
    StatusEnum,
    ListedDocument,
    DocumentDownloadResponse,
)
from app.tenant_config import TenantSettings
from app.utils.logger import get_logger

logger = get_logger(__name__)


class EntrMailboxRetrievalHandler(RetrievalHandler):
    """
    RetrievalHandler that pulls mail from the shared intake mailbox in Gmail.

    Each tenant only sees messages that have that tenant's label.
    We expose each email as a 'document'. When requested, we return
    the ENTIRE raw RFC822 email (.eml), base64 encoded.

    Auth model:
    - OAuth2 client_id/client_secret + refresh_token
      (refresh_token was granted by entr-processor@entr.solutions).
    - We exchange refresh_token -> access_token on demand.

    Isolation model:
    - TenantSettings.TENANT_NAME maps to a Gmail label.
    - We only list/read messages carrying that label.
    """

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

    def __init__(self, settings: TenantSettings, label_override: str = None):
        super().__init__(settings)
        self._label_override = label_override
        logger.info(f"Initialized EntrMailboxRetrievalHandler for tenant: {self.settings.TENANT_NAME}"
                     + (f" (label override: {label_override})" if label_override else ""))
        self._label_id_cache: Optional[str] = None

    # -------------------------
    # Internal helpers
    # -------------------------

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

        When a label_override is provided (for multi-division setups),
        return that instead of the default tenant-based label.
        """
        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(f"Failed to refresh access token: {resp.status_code} {resp.text}")
            raise RuntimeError("OAuth token refresh failed")
        token_json = resp.json()
        return token_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.
        Cache after first lookup.
        """
        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(f"Failed to list labels: {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(f"Gmail label '{desired_label}' not found")
        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]:
        """
        Returns all message stubs { "id": "...", "threadId": "..." } for mail
        with this label, using pagination to fetch all results.

        Args:
            date_from: If provided, only return messages received on or after this date
        """
        from urllib.parse import quote

        base_url = f"{self.GMAIL_API_BASE}/users/me/messages?labelIds={label_id}&maxResults=500"

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

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

        # Paginate through all results
        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(f"Failed to list messages: {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(f"Fetching next page, {len(all_messages)} messages so far...")

        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).
        Gmail returns base64url-encoded string of the entire .eml.
        We'll convert that to binary bytes.
        """
        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(f"Failed to get raw message {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).
        # Convert base64url -> standard base64 -> bytes.
        b64std = raw_b64url.replace("-", "+").replace("_", "/")
        # pad with '=' to length %4 == 0
        padding_needed = (4 - (len(b64std) % 4)) % 4
        b64std += "=" * padding_needed
        raw_bytes = base64.b64decode(b64std)

        return raw_bytes

    # -------------------------
    # Public API
    # -------------------------

    async def list_documents(self, payload: RetrievalPayload) -> RetrievalResponse:
        """
        Return recent emails for this tenant's label.
        Each email becomes a ListedDocument with:
        - document_identifier = Gmail message ID
        - filename = Message ID with .eml extension

        Date filtering is done server-side by Gmail for efficiency.
        """
        logger.info(
            f"EntrMailboxRetrievalHandler.list_documents "
            f"tenant={self.settings.TENANT_NAME} source={payload.source_identifier} "
            f"date_from={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)

            # Gmail filters by date server-side, with pagination for high volume
            msgs = await self._list_messages_for_label(
                client, access_token, label_id,
                date_from=payload.date_from
            )

            logger.info(f"Gmail returned {len(msgs)} emails matching filters")

            # Build document list directly from message IDs
            # Subject extraction happens later when Core processes the .eml
            docs: List[ListedDocument] = [
                ListedDocument(
                    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 this Gmail message, base64-encoded.
        """
        logger.info(
            f"EntrMailboxRetrievalHandler.get_document "
            f"tenant={self.settings.TENANT_NAME} message_id={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)

        # The response model expects base64 of the binary document.
        eml_b64 = base64.b64encode(raw_bytes).decode("utf-8")

        # Give downstream a stable filename that ends in .eml
        filename = f"{document_identifier}.eml"

        return DocumentDownloadResponse(
            document_identifier=filename,         # human-ish identifier
            adapter_document_id=document_identifier,  # original gmail message id
            content_b64=eml_b64,
        )
