task_manager_agent/data_fetcher.py

"""
data_fetcher.py — Task Manager Agent
=====================================
Pulls raw data from Gmail, Google Calendar, and existing task stores.
"""

import os
import base64
import logging
from datetime import datetime, timedelta, timezone
from typing import Optional

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

logger = logging.getLogger("TaskManagerAgent.DataFetcher")

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

CREDENTIALS_FILE = os.getenv("GOOGLE_CREDENTIALS_FILE")
TOKEN_FILE = os.getenv("GOOGLE_TOKEN_FILE")

# Labels that signal actionable emails (expand as needed)
ACTIONABLE_LABELS = {"INBOX", "IMPORTANT", "STARRED"}


# ── Google auth ───────────────────────────────────────────────────────────────

def _get_google_creds() -> Credentials:
    creds = None
    if os.path.exists(TOKEN_FILE):
        creds = Credentials.from_authorized_user_file(TOKEN_FILE, SCOPES)
    if not creds or not creds.valid:
        if creds and creds.expired and creds.refresh_token:
            creds.refresh(Request())
        else:
            flow = InstalledAppFlow.from_client_secrets_file(CREDENTIALS_FILE, SCOPES)
            creds = flow.run_local_server(port=0)
        with open(TOKEN_FILE, "w") as f:
            f.write(creds.to_json())
    return creds


def _gmail_service():
    return build("gmail", "v1", credentials=_get_google_creds())


def _calendar_service():
    return build("calendar", "v3", credentials=_get_google_creds())


# ── Gmail ─────────────────────────────────────────────────────────────────────

def fetch_recent_emails(
    max_results: int = 20,
    hours_back: int = 24,
    since_id: Optional[str] = None,
) -> list[dict]:
    """
    Returns list of dicts with keys:
      id, subject, sender, date, snippet, body (plain text, truncated)
    """
    try:
        service = _gmail_service()
        after_ts = int((datetime.now(timezone.utc) - timedelta(hours=hours_back)).timestamp())
        query = f"after:{after_ts} -category:promotions -category:social"

        result = service.users().messages().list(
            userId="me", q=query, maxResults=max_results
        ).execute()
        messages = result.get("messages", [])

        emails = []
        for msg_ref in messages:
            if since_id and msg_ref["id"] == since_id:
                break
            try:
                msg = service.users().messages().get(
                    userId="me", id=msg_ref["id"], format="full"
                ).execute()
                emails.append(_parse_email(msg))
            except Exception as e:
                logger.warning(f"Could not fetch message {msg_ref['id']}: {e}")

        logger.debug(f"Fetched {len(emails)} emails")
        return emails

    except Exception as e:
        logger.error(f"Gmail fetch error: {e}", exc_info=True)
        return []


def _parse_email(msg: dict) -> dict:
    headers = {h["name"]: h["value"] for h in msg["payload"].get("headers", [])}
    body = _extract_body(msg["payload"])
    return {
        "id": msg["id"],
        "thread_id": msg.get("threadId", ""),
        "subject": headers.get("Subject", "(no subject)"),
        "sender": headers.get("From", ""),
        "to": headers.get("To", ""),
        "date": headers.get("Date", ""),
        "snippet": msg.get("snippet", ""),
        "body": body[:3000],   # cap at 3k chars
        "labels": msg.get("labelIds", []),
    }


def _extract_body(payload: dict) -> str:
    """Recursively extracts plain-text body from MIME payload."""
    if payload.get("mimeType") == "text/plain":
        data = payload.get("body", {}).get("data", "")
        if data:
            return base64.urlsafe_b64decode(data).decode("utf-8", errors="replace")
    for part in payload.get("parts", []):
        text = _extract_body(part)
        if text:
            return text
    return ""


# ── Google Calendar ───────────────────────────────────────────────────────────

def fetch_calendar_meetings(
    days_back: int = 1,
    days_ahead: int = 0,
    calendar_id: str = "primary",
) -> list[dict]:
    """
    Returns meetings from [now - days_back, now + days_ahead].
    Each dict: id, title, start, end, attendees, description, location
    """
    try:
        service = _calendar_service()
        now = datetime.now(timezone.utc)
        time_min = (now - timedelta(days=days_back)).isoformat()
        time_max = (now + timedelta(days=days_ahead + 1)).isoformat()

        events_result = service.events().list(
            calendarId=calendar_id,
            timeMin=time_min,
            timeMax=time_max,
            maxResults=50,
            singleEvents=True,
            orderBy="startTime",
        ).execute()
        events = events_result.get("items", [])

        meetings = []
        for e in events:
            if e.get("status") == "cancelled":
                continue
            meetings.append({
                "id": e.get("id", ""),
                "title": e.get("summary", "(untitled)"),
                "start": e.get("start", {}).get("dateTime", e.get("start", {}).get("date", "")),
                "end": e.get("end", {}).get("dateTime", e.get("end", {}).get("date", "")),
                "attendees": [
                    a.get("email", "") for a in e.get("attendees", [])
                ],
                "description": e.get("description", "")[:2000],
                "location": e.get("location", ""),
            })

        logger.debug(f"Fetched {len(meetings)} meetings")
        return meetings

    except Exception as e:
        logger.error(f"Calendar fetch error: {e}", exc_info=True)
        return []


# ── Existing tasks (local store read) ─────────────────────────────────────────

def fetch_existing_tasks() -> list[dict]:
    """Loads tasks from local task store + Notion for deduplication context."""
    tasks = []

    # local
    local_path = os.getenv("LOCAL_TASKS_FILE", "tasks.json")
    if os.path.exists(local_path):
        import json
        try:
            with open(local_path) as f:
                tasks.extend(json.load(f))
            logger.debug(f"Loaded {len(tasks)} local tasks")
        except Exception as e:
            logger.warning(f"Could not load local tasks: {e}")

    return tasks


def fetch_notion_tasks() -> list[dict]:
    """Fetches open tasks from Notion DB."""
    import requests

    notion_token = os.getenv("NOTION_TOKEN")
    db_id = os.getenv("NOTION_TASKS_DB_ID")
    if not notion_token or not db_id:
        logger.debug("Notion credentials not set — skipping Notion task fetch")
        return []

    try:
        headers = {
            "Authorization": f"Bearer {notion_token}",
            "Notion-Version": "2022-06-28",
            "Content-Type": "application/json",
        }
        payload = {
            "filter": {
                "property": "Status",
                "select": {"does_not_equal": "Done"},
            }
        }
        resp = requests.post(
            f"https://api.notion.com/v1/databases/{db_id}/query",
            headers=headers, json=payload, timeout=10,
        )
        resp.raise_for_status()
        results = resp.json().get("results", [])

        tasks = []
        for page in results:
            props = page.get("properties", {})
            title_prop = props.get("Name", {}).get("title", [])
            title = title_prop[0]["plain_text"] if title_prop else "(no title)"
            tasks.append({
                "id": page["id"],
                "title": title,
                "source": "notion",
                "url": page.get("url", ""),
            })
        logger.debug(f"Fetched {len(tasks)} tasks from Notion")
        return tasks

    except Exception as e:
        logger.error(f"Notion fetch error: {e}", exc_info=True)
        return []