email_briefing.py

"""
Email briefing — "Connect your school Gmail and the Companion greets you
with a spoken summary of what's new" feature.

Flow
----
1. Student clicks "Connect Google account" in Settings.
   Frontend calls POST /api/email/connect (Firebase-authed) and gets back a
   Google OAuth URL bound to that student via a signed `state` parameter.
2. Google redirects to GET /api/email/oauth/callback with a code. We exchange
   it for a refresh token (gmail.readonly scope only), verify the Google
   account matches the student's ALU email, store it, and bounce the browser
   back to the frontend Settings page.
3. On chat open the frontend calls GET /api/email/briefing. We mint a fresh
   access token from the stored refresh token, pull recent inbox messages via
   the Gmail REST API, filter/summarize them with the existing LLM chain
   (Claude -> NVIDIA -> Groq), and return a short spoken-style briefing that
   the frontend reads aloud with the Web Speech API.

Configuration (HF Space secrets)
--------------------------------
    GOOGLE_OAUTH_CLIENT_ID      OAuth client (Web application type)
    GOOGLE_OAUTH_CLIENT_SECRET
    BACKEND_PUBLIC_URL          e.g. https://studentcompanion-alu-chatbot.hf.space
                                (used to build the redirect URI Google calls)
    FRONTEND_URL                e.g. https://chat.studentcompanionai.rw
    EMAIL_TOKEN_KEY             optional Fernet key; if set, refresh tokens are
                                encrypted at rest in Aurora

Storage: Aurora table `email_accounts` when the data layer is configured,
otherwise an in-process dict (fine for local dev; lost on restart).
"""

import os
import json
import time
import hmac
import base64
import hashlib
import logging
from urllib.parse import urlencode
from typing import Any, Dict, List, Optional

import requests
from fastapi import APIRouter, Depends, HTTPException, Query
from fastapi.responses import RedirectResponse
from pydantic import BaseModel

import db
from auth import require_user

logger = logging.getLogger("email_briefing")

router = APIRouter(tags=["email-briefing"])

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
CLIENT_ID = os.environ.get("GOOGLE_OAUTH_CLIENT_ID", "").strip()
CLIENT_SECRET = os.environ.get("GOOGLE_OAUTH_CLIENT_SECRET", "").strip()
BACKEND_PUBLIC_URL = os.environ.get(
    "BACKEND_PUBLIC_URL", "https://studentcompanion-alu-chatbot.hf.space"
).rstrip("/")
FRONTEND_URL = os.environ.get(
    "FRONTEND_URL", "https://chat.studentcompanionai.rw"
).rstrip("/")
REDIRECT_URI = f"{BACKEND_PUBLIC_URL}/api/email/oauth/callback"

GMAIL_SCOPE = "https://www.googleapis.com/auth/gmail.readonly"
OAUTH_AUTH_URL = "https://accounts.google.com/o/oauth2/v2/auth"
OAUTH_TOKEN_URL = "https://oauth2.googleapis.com/token"
GMAIL_API = "https://gmail.googleapis.com/gmail/v1/users/me"

STATE_TTL_SECONDS = 600  # OAuth round-trip must complete within 10 minutes

DEFAULT_PREFS: Dict[str, Any] = {
    # What the briefing should mention
    "assignments": True,       # assignments & due dates
    "classes": True,           # upcoming classes / schedule changes
    "congratulations": True,   # achievements & congratulations
    "announcements": True,     # school announcements & events
    "other": False,            # everything else in the inbox
    # How it should behave
    "voice_enabled": True,     # read the briefing aloud
    "auto_play": True,         # speak automatically when the chat opens
    "lookback_hours": 48,      # how far back to scan the inbox
}

PREF_CATEGORY_LABELS = {
    "assignments": "assignments and due dates",
    "classes": "upcoming classes and schedule changes",
    "congratulations": "congratulations and achievements",
    "announcements": "school announcements and events",
    "other": "any other notable emails",
}


def oauth_configured() -> bool:
    return bool(CLIENT_ID and CLIENT_SECRET)


# ---------------------------------------------------------------------------
# Signed OAuth state — binds the Google callback to the Firebase user who
# started the flow, without server-side session storage.
# ---------------------------------------------------------------------------
def _state_secret() -> bytes:
    return hashlib.sha256(
        ("email-briefing-state:" + CLIENT_SECRET).encode()
    ).digest()


def _sign(payload: bytes) -> str:
    return hmac.new(_state_secret(), payload, hashlib.sha256).hexdigest()[:32]


def make_state(uid: str, email: str) -> str:
    payload = base64.urlsafe_b64encode(
        json.dumps({"uid": uid, "email": email, "exp": int(time.time()) + STATE_TTL_SECONDS}).encode()
    ).decode()
    return f"{payload}.{_sign(payload.encode())}"


def parse_state(state: str) -> Optional[Dict[str, Any]]:
    try:
        payload, sig = state.rsplit(".", 1)
        if not hmac.compare_digest(sig, _sign(payload.encode())):
            return None
        data = json.loads(base64.urlsafe_b64decode(payload.encode()))
        if data.get("exp", 0) < time.time():
            return None
        return data
    except Exception:
        return None


# ---------------------------------------------------------------------------
# Refresh-token storage — Aurora when available, else in-process memory.
# Tokens are optionally encrypted with EMAIL_TOKEN_KEY (Fernet).
# ---------------------------------------------------------------------------
_memory_store: Dict[str, Dict[str, Any]] = {}
_table_ready = False

_fernet = None
_key = os.environ.get("EMAIL_TOKEN_KEY", "").strip()
if _key:
    try:
        from cryptography.fernet import Fernet

        _fernet = Fernet(_key.encode())
    except Exception as e:  # bad key or lib missing — refuse silently-degraded crypto
        logger.error("[email] EMAIL_TOKEN_KEY set but unusable (%s); tokens stored raw", e)


def _seal(token: str) -> str:
    return _fernet.encrypt(token.encode()).decode() if _fernet else token


def _unseal(token: str) -> str:
    if _fernet:
        try:
            return _fernet.decrypt(token.encode()).decode()
        except Exception:
            return token  # stored before the key existed
    return token


def _ensure_table(conn) -> None:
    global _table_ready
    if _table_ready:
        return
    conn.execute(
        """
        CREATE TABLE IF NOT EXISTS email_accounts (
            user_email    TEXT PRIMARY KEY,
            google_email  TEXT NOT NULL,
            refresh_token TEXT NOT NULL,
            prefs         JSONB NOT NULL DEFAULT '{}'::jsonb,
            connected_at  TIMESTAMPTZ NOT NULL DEFAULT now()
        )
        """
    )
    _table_ready = True


def _store_get(user_email: str) -> Optional[Dict[str, Any]]:
    pool = db.get_pool()
    if pool is None:
        return _memory_store.get(user_email)
    try:
        with pool.connection() as conn:
            _ensure_table(conn)
            row = conn.execute(
                "SELECT google_email, refresh_token, prefs FROM email_accounts WHERE user_email = %s",
                (user_email,),
            ).fetchone()
        if not row:
            return None
        prefs = row["prefs"] if isinstance(row["prefs"], dict) else json.loads(row["prefs"] or "{}")
        return {
            "google_email": row["google_email"],
            "refresh_token": _unseal(row["refresh_token"]),
            "prefs": prefs,
        }
    except Exception as e:
        logger.error("[email] store_get failed: %s", e)
        return _memory_store.get(user_email)


def _store_put(user_email: str, google_email: str, refresh_token: str, prefs: Dict[str, Any]) -> None:
    pool = db.get_pool()
    if pool is None:
        _memory_store[user_email] = {
            "google_email": google_email,
            "refresh_token": refresh_token,
            "prefs": prefs,
        }
        return
    try:
        with pool.connection() as conn:
            _ensure_table(conn)
            conn.execute(
                """
                INSERT INTO email_accounts (user_email, google_email, refresh_token, prefs)
                VALUES (%s, %s, %s, %s)
                ON CONFLICT (user_email) DO UPDATE SET
                    google_email = EXCLUDED.google_email,
                    refresh_token = EXCLUDED.refresh_token,
                    prefs = EXCLUDED.prefs,
                    connected_at = now()
                """,
                (user_email, google_email, _seal(refresh_token), json.dumps(prefs)),
            )
    except Exception as e:
        logger.error("[email] store_put failed, keeping in memory: %s", e)
        _memory_store[user_email] = {
            "google_email": google_email,
            "refresh_token": refresh_token,
            "prefs": prefs,
        }


def _store_update_prefs(user_email: str, prefs: Dict[str, Any]) -> bool:
    account = _store_get(user_email)
    if account is None:
        return False
    _store_put(user_email, account["google_email"], account["refresh_token"], prefs)
    return True


def _store_delete(user_email: str) -> None:
    _memory_store.pop(user_email, None)
    pool = db.get_pool()
    if pool is None:
        return
    try:
        with pool.connection() as conn:
            _ensure_table(conn)
            conn.execute("DELETE FROM email_accounts WHERE user_email = %s", (user_email,))
    except Exception as e:
        logger.error("[email] store_delete failed: %s", e)


# ---------------------------------------------------------------------------
# Google API helpers (plain REST — no extra client libraries)
# ---------------------------------------------------------------------------
def _exchange_code(code: str) -> Dict[str, Any]:
    resp = requests.post(
        OAUTH_TOKEN_URL,
        data={
            "code": code,
            "client_id": CLIENT_ID,
            "client_secret": CLIENT_SECRET,
            "redirect_uri": REDIRECT_URI,
            "grant_type": "authorization_code",
        },
        timeout=15,
    )
    resp.raise_for_status()
    return resp.json()


def _refresh_access_token(refresh_token: str) -> Optional[str]:
    """Mint a short-lived access token. Returns None if the grant was revoked."""
    resp = requests.post(
        OAUTH_TOKEN_URL,
        data={
            "refresh_token": refresh_token,
            "client_id": CLIENT_ID,
            "client_secret": CLIENT_SECRET,
            "grant_type": "refresh_token",
        },
        timeout=15,
    )
    if resp.status_code != 200:
        logger.warning("[email] token refresh failed (%s): %s", resp.status_code, resp.text[:200])
        return None
    return resp.json().get("access_token")


def _id_token_email(id_token: str) -> str:
    """Extract the email claim from a Google id_token (payload only — the token
    just arrived over TLS directly from Google's token endpoint)."""
    try:
        payload = id_token.split(".")[1]
        payload += "=" * (-len(payload) % 4)
        return (json.loads(base64.urlsafe_b64decode(payload)).get("email") or "").lower()
    except Exception:
        return ""


def _fetch_recent_emails(access_token: str, lookback_hours: int) -> List[Dict[str, str]]:
    """Return [{from, subject, date, snippet}] for recent inbox messages."""
    headers = {"Authorization": f"Bearer {access_token}"}
    days = max(1, round(lookback_hours / 24))
    listing = requests.get(
        f"{GMAIL_API}/messages",
        params={"q": f"in:inbox newer_than:{days}d", "maxResults": 20},
        headers=headers,
        timeout=15,
    )
    listing.raise_for_status()
    ids = [m["id"] for m in listing.json().get("messages", [])]

    emails: List[Dict[str, str]] = []
    for mid in ids:
        try:
            msg = requests.get(
                f"{GMAIL_API}/messages/{mid}",
                params={
                    "format": "metadata",
                    "metadataHeaders": ["Subject", "From", "Date"],
                },
                headers=headers,
                timeout=15,
            )
            msg.raise_for_status()
            data = msg.json()
            hdrs = {
                h["name"].lower(): h["value"]
                for h in data.get("payload", {}).get("headers", [])
            }
            emails.append(
                {
                    "from": hdrs.get("from", "Unknown sender"),
                    "subject": hdrs.get("subject", "(no subject)"),
                    "date": hdrs.get("date", ""),
                    "snippet": (data.get("snippet") or "")[:200],
                }
            )
        except Exception as e:
            logger.warning("[email] failed to fetch message %s: %s", mid, e)
    return emails


# ---------------------------------------------------------------------------
# Briefing composition
# ---------------------------------------------------------------------------
def _first_name(user: Dict[str, Any]) -> str:
    name = (user.get("name") or "").strip()
    if name:
        return name.split()[0]
    email = user.get("email") or "there"
    return email.split("@")[0].split(".")[0].capitalize()


def _briefing_prompt(first_name: str, prefs: Dict[str, Any], emails: List[Dict[str, str]]) -> str:
    wanted = [
        label
        for key, label in PREF_CATEGORY_LABELS.items()
        if prefs.get(key, DEFAULT_PREFS[key])
    ]
    digest = "\n".join(
        f"- From: {e['from']} | Subject: {e['subject']} | Preview: {e['snippet']}"
        for e in emails
    )
    return (
        f"You are the ALU Student Companion greeting {first_name} as they open the app. "
        f"Below are their recent inbox emails. Write a short spoken greeting that starts "
        f"with 'Hi {first_name}' and summarizes ONLY emails in these categories: "
        f"{', '.join(wanted) or 'none'}. Ignore everything else (promotions, newsletters, "
        f"automated notifications) unless it fits a selected category. "
        f"If a congratulations email is there, congratulate them warmly yourself. "
        f"Mention senders by name, mention due dates and times exactly as written. "
        f"If nothing matches the selected categories, say their inbox has nothing "
        f"important right now in one friendly sentence. "
        f"Rules: plain conversational text that sounds natural read aloud, no markdown, "
        f"no emojis, no bullet points, at most 120 words.\n\n"
        f"RECENT EMAILS:\n{digest if digest else '(inbox is empty)'}"
    )


def _fallback_briefing(first_name: str, emails: List[Dict[str, str]]) -> str:
    if not emails:
        return f"Hi {first_name}. Your inbox is all clear — no new emails in the last couple of days."
    top = emails[:3]
    parts = [f"Hi {first_name}. You have {len(emails)} recent email{'s' if len(emails) != 1 else ''}."]
    for e in top:
        sender = e["from"].split("<")[0].strip().strip('"') or "someone"
        parts.append(f"{sender} wrote about {e['subject']}.")
    return " ".join(parts)


def _llm_briefing(prompt: str) -> Optional[Dict[str, str]]:
    """Run the prompt through the same engine chain the chat uses."""
    try:
        from claude_engine import claude_engine

        if claude_engine.enabled:
            answer = claude_engine.generate_response(query=prompt, context_docs=None)
            if answer:
                return {"text": answer.strip(), "engine": "claude"}
    except Exception as e:
        logger.warning("[email] claude briefing failed: %s", e)
    try:
        from nvidia_fallback import nvidia_fallback

        if nvidia_fallback and nvidia_fallback.enabled:
            answer = nvidia_fallback.generate_response(prompt, "")
            if answer:
                return {"text": answer.strip(), "engine": "nvidia"}
    except Exception as e:
        logger.warning("[email] nvidia briefing failed: %s", e)
    try:
        from main import groq_fallback

        if groq_fallback and groq_fallback.enabled:
            answer = groq_fallback.generate_response(prompt, "")
            if answer:
                return {"text": answer.strip(), "engine": "groq"}
    except Exception as e:
        logger.warning("[email] groq briefing failed: %s", e)
    return None


# Briefings are cached briefly so reopening the chat doesn't re-scan Gmail
# and re-run the LLM every time.
_briefing_cache: Dict[str, Dict[str, Any]] = {}
BRIEFING_CACHE_SECONDS = 600


# ---------------------------------------------------------------------------
# Routes (mounted under /api/email in app.py)
# ---------------------------------------------------------------------------
class PrefsBody(BaseModel):
    prefs: Dict[str, Any]


@router.get("/status")
async def email_status(user: dict = Depends(require_user)):
    """Is the feature configured server-side, and has this student connected?"""
    account = _store_get((user.get("email") or "").lower())
    prefs = {**DEFAULT_PREFS, **(account.get("prefs") or {})} if account else dict(DEFAULT_PREFS)
    return {
        "configured": oauth_configured(),
        "connected": account is not None,
        "google_email": account["google_email"] if account else None,
        "prefs": prefs,
    }


@router.post("/connect")
async def email_connect(user: dict = Depends(require_user)):
    """Return the Google OAuth URL for this student to approve Gmail access."""
    if not oauth_configured():
        raise HTTPException(
            503,
            "Email briefing is not configured on the server "
            "(GOOGLE_OAUTH_CLIENT_ID / GOOGLE_OAUTH_CLIENT_SECRET missing).",
        )
    email = (user.get("email") or "").lower()
    params = {
        "client_id": CLIENT_ID,
        "redirect_uri": REDIRECT_URI,
        "response_type": "code",
        "scope": f"openid email {GMAIL_SCOPE}",
        "access_type": "offline",
        "prompt": "consent",
        "login_hint": email,
        "state": make_state(user.get("uid", ""), email),
    }
    return {"auth_url": f"{OAUTH_AUTH_URL}?{urlencode(params)}"}


@router.get("/oauth/callback")
async def email_oauth_callback(
    code: str = Query(default=""),
    state: str = Query(default=""),
    error: str = Query(default=""),
):
    """Google redirects here. No Firebase auth — identity comes from `state`."""
    settings_url = f"{FRONTEND_URL}/settings"
    if error or not code:
        return RedirectResponse(f"{settings_url}?email_briefing=denied")

    data = parse_state(state)
    if not data:
        return RedirectResponse(f"{settings_url}?email_briefing=expired")

    try:
        tokens = _exchange_code(code)
    except Exception as e:
        logger.error("[email] code exchange failed: %s", e)
        return RedirectResponse(f"{settings_url}?email_briefing=error")

    refresh_token = tokens.get("refresh_token")
    if not refresh_token:
        # Google omits it if the app was already authorized without prompt=consent
        return RedirectResponse(f"{settings_url}?email_briefing=error")

    google_email = _id_token_email(tokens.get("id_token", ""))
    if google_email and google_email != data["email"]:
        # The student authorized a different Google account than their ALU login.
        return RedirectResponse(f"{settings_url}?email_briefing=mismatch")

    existing = _store_get(data["email"])
    prefs = (existing or {}).get("prefs") or dict(DEFAULT_PREFS)
    _store_put(data["email"], google_email or data["email"], refresh_token, prefs)
    logger.info("[email] connected mailbox for %s", data["email"])
    return RedirectResponse(f"{settings_url}?email_briefing=connected")


@router.post("/disconnect")
async def email_disconnect(user: dict = Depends(require_user)):
    email = (user.get("email") or "").lower()
    account = _store_get(email)
    if account:
        # Best-effort revoke at Google so the grant doesn't linger.
        try:
            requests.post(
                "https://oauth2.googleapis.com/revoke",
                params={"token": account["refresh_token"]},
                timeout=10,
            )
        except Exception:
            pass
    _store_delete(email)
    _briefing_cache.pop(email, None)
    return {"connected": False}


@router.put("/preferences")
async def email_preferences(body: PrefsBody, user: dict = Depends(require_user)):
    email = (user.get("email") or "").lower()
    clean = {k: v for k, v in body.prefs.items() if k in DEFAULT_PREFS}
    if not _store_update_prefs(email, {**DEFAULT_PREFS, **clean}):
        raise HTTPException(404, "No connected email account")
    _briefing_cache.pop(email, None)  # prefs changed — regenerate next time
    return {"prefs": {**DEFAULT_PREFS, **clean}}


@router.get("/briefing")
async def email_briefing(
    refresh: bool = Query(default=False),
    user: dict = Depends(require_user),
):
    email = (user.get("email") or "").lower()
    account = _store_get(email)
    if account is None:
        return {"connected": False, "briefing": None}

    cached = _briefing_cache.get(email)
    if cached and not refresh and time.time() - cached["at"] < BRIEFING_CACHE_SECONDS:
        return {**cached["payload"], "cached": True}

    access_token = _refresh_access_token(account["refresh_token"])
    if access_token is None:
        # Grant revoked/expired — force a clean reconnect instead of erroring forever.
        _store_delete(email)
        return {"connected": False, "briefing": None, "reason": "reconnect_required"}

    prefs = {**DEFAULT_PREFS, **(account.get("prefs") or {})}
    try:
        emails = _fetch_recent_emails(access_token, int(prefs.get("lookback_hours", 48)))
    except Exception as e:
        logger.error("[email] gmail fetch failed: %s", e)
        raise HTTPException(502, "Could not read the inbox from Gmail")

    first_name = _first_name(user)
    result = _llm_briefing(_briefing_prompt(first_name, prefs, emails))
    if result is None:
        result = {"text": _fallback_briefing(first_name, emails), "engine": "template"}

    payload = {
        "connected": True,
        "briefing": result["text"],
        "engine": result["engine"],
        "email_count": len(emails),
        "prefs": prefs,
        "generated_at": int(time.time()),
    }
    _briefing_cache[email] = {"at": time.time(), "payload": payload}
    return payload