simbarag/utils/google_service.py

"""Google Calendar and Gmail service using the Python SDK."""

import asyncio
import base64
import os
from typing import Any

from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from googleapiclient.discovery import build

SCOPES = [
    "https://www.googleapis.com/auth/calendar.events",
    "https://www.googleapis.com/auth/gmail.readonly",
]


class GoogleService:
    def __init__(self):
        token_path = os.getenv("GOOGLE_TOKEN_FILE", "/app/config/google-token.json")
        if not os.path.exists(token_path):
            raise FileNotFoundError(f"Google token file not found: {token_path}")

        self._token_path = token_path
        creds = Credentials.from_authorized_user_file(token_path, SCOPES)
        if creds.expired and creds.refresh_token:
            creds.refresh(Request())
            with open(token_path, "w") as f:
                f.write(creds.to_json())
        self._creds = creds

        self._calendar = build("calendar", "v3", credentials=creds)
        self._gmail = build("gmail", "v1", credentials=creds)

    async def get_calendar_events(
        self,
        time_min: str,
        time_max: str,
        calendar_id: str = "primary",
        max_results: int = 20,
    ) -> list[dict[str, Any]]:
        def _fetch():
            result = (
                self._calendar.events()
                .list(
                    calendarId=calendar_id,
                    timeMin=time_min,
                    timeMax=time_max,
                    maxResults=max_results,
                    singleEvents=True,
                    orderBy="startTime",
                )
                .execute()
            )
            return result.get("items", [])

        return await asyncio.to_thread(_fetch)

    async def create_calendar_event(
        self,
        summary: str,
        start_time: str,
        end_time: str,
        description: str = "",
        location: str = "",
        calendar_id: str = "primary",
    ) -> dict[str, Any]:
        def _create():
            body: dict[str, Any] = {
                "summary": summary,
                "start": {"dateTime": start_time},
                "end": {"dateTime": end_time},
            }
            if description:
                body["description"] = description
            if location:
                body["location"] = location
            return (
                self._calendar.events()
                .insert(calendarId=calendar_id, body=body)
                .execute()
            )

        return await asyncio.to_thread(_create)

    async def search_emails(
        self,
        query: str,
        max_results: int = 10,
    ) -> list[dict[str, Any]]:
        def _fetch():
            result = (
                self._gmail.users()
                .messages()
                .list(userId="me", q=query, maxResults=max_results)
                .execute()
            )
            messages = result.get("messages", [])

            emails = []
            for msg in messages:
                msg_data = (
                    self._gmail.users()
                    .messages()
                    .get(
                        userId="me",
                        id=msg["id"],
                        format="metadata",
                        metadataHeaders=["From", "To", "Subject", "Date"],
                    )
                    .execute()
                )
                headers = {
                    h["name"]: h["value"]
                    for h in msg_data.get("payload", {}).get("headers", [])
                }
                emails.append(
                    {
                        "id": msg_data["id"],
                        "from": headers.get("From", ""),
                        "to": headers.get("To", ""),
                        "subject": headers.get("Subject", ""),
                        "date": headers.get("Date", ""),
                        "snippet": msg_data.get("snippet", ""),
                    }
                )
            return emails

        return await asyncio.to_thread(_fetch)

    async def get_email(self, message_id: str) -> dict[str, Any]:
        def _fetch():
            msg_data = (
                self._gmail.users()
                .messages()
                .get(userId="me", id=message_id, format="full")
                .execute()
            )
            headers = {
                h["name"]: h["value"]
                for h in msg_data.get("payload", {}).get("headers", [])
            }

            body = _extract_body(msg_data.get("payload", {}))

            return {
                "id": msg_data["id"],
                "from": headers.get("From", ""),
                "to": headers.get("To", ""),
                "subject": headers.get("Subject", ""),
                "date": headers.get("Date", ""),
                "body": body[:4000],
                "snippet": msg_data.get("snippet", ""),
            }

        return await asyncio.to_thread(_fetch)


def _extract_body(payload: dict[str, Any]) -> str:
    mime = payload.get("mimeType", "")
    data = payload.get("body", {}).get("data", "")

    if mime == "text/plain" and data:
        return base64.urlsafe_b64decode(data).decode("utf-8", errors="replace")

    if mime == "text/html" and data:
        import html2text

        html = base64.urlsafe_b64decode(data).decode("utf-8", errors="replace")
        return html2text.html2text(html)

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

    return ""