routes/auth.py

import os
import time
from datetime import timedelta
from functools import lru_cache

from fastapi import HTTPException, Security
from fastapi.security import HTTPAuthorizationCredentials, HTTPBearer
import jwt
from jwt import PyJWKClient

_security = HTTPBearer()

# profiles.status lookups are cached briefly so each request doesn't cost a
# DB round-trip. Admin approve/reject calls invalidate_status_cache() so the
# change takes effect immediately.
_STATUS_TTL_SECONDS = 60
_status_cache: dict[str, tuple[str, float]] = {}
# profiles."Source access" is cached the same way and gates which data-source
# fact tables the SQL agent may query for a user.
_source_access_cache: dict[str, tuple[str | None, float]] = {}


@lru_cache(maxsize=1)
def _service_client():
    from supabase import create_client

    url = os.getenv("SUPABASE_URL", "")
    key = os.getenv("SUPABASE_SERVICE_ROLE_KEY", "")
    if not url or not key:
        raise RuntimeError("SUPABASE_URL and SUPABASE_SERVICE_ROLE_KEY must be set in environment")
    return create_client(url, key)


@lru_cache(maxsize=1)
def _jwks_client() -> PyJWKClient:
    url = os.getenv("SUPABASE_URL", "").rstrip("/")
    if not url:
        raise RuntimeError("SUPABASE_URL is not set in environment.")
    return PyJWKClient(f"{url}/auth/v1/.well-known/jwks.json")


def _decode(token: str) -> dict:
    """Verify a Supabase access token, supporting both signing modes.

    Asymmetric tokens (RS256/ES256) carry a `kid` and are verified against the
    project's published JWKS. Legacy symmetric tokens (HS256) have no `kid` and
    are verified with the shared SUPABASE_JWT_SECRET. Supporting both lets the
    backend keep working across a JWKS migration, regardless of which signing
    key the project currently issues tokens with.
    """
    header = jwt.get_unverified_header(token)
    common = dict(
        audience="authenticated",
        leeway=timedelta(seconds=60),  # tolerate up to 60 s of clock skew
        # `iat` is informational, not a not-before claim (that's `nbf`). Newer
        # PyJWT rejects tokens whose `iat` is in the future, which fires on even
        # modest server/issuer clock skew. Skip that check; signature, `exp`,
        # and `aud` are still fully verified.
        options={"verify_iat": False},
    )
    if header.get("kid") and header.get("alg") in ("RS256", "ES256"):
        signing_key = _jwks_client().get_signing_key_from_jwt(token)
        return jwt.decode(token, signing_key.key, algorithms=["RS256", "ES256"], **common)

    secret = os.getenv("SUPABASE_JWT_SECRET", "")
    if not secret:
        raise HTTPException(
            status_code=401,
            detail="HS256 token received but SUPABASE_JWT_SECRET is not configured",
        )
    return jwt.decode(token, secret, algorithms=["HS256"], **common)


def _decode_verified(credentials: HTTPAuthorizationCredentials) -> dict:
    """Verify a bearer token and return its claims, mapping JWT errors to 401s."""
    try:
        return _decode(credentials.credentials)
    except jwt.ExpiredSignatureError:
        raise HTTPException(status_code=401, detail="Token expired")
    except jwt.InvalidTokenError as exc:
        raise HTTPException(status_code=401, detail=f"Invalid token: {exc}")


def verify_token(credentials: HTTPAuthorizationCredentials = Security(_security)) -> str:
    payload = _decode_verified(credentials)
    user_id: str = payload.get("sub", "")
    if not user_id:
        raise HTTPException(status_code=401, detail="Token missing subject claim")
    return user_id


def verify_token_email(credentials: HTTPAuthorizationCredentials = Security(_security)) -> str:
    """Like verify_token, but returns the authenticated user's login email (lower-cased).

    Supabase access tokens carry an `email` claim, so this needs no DB lookup. Used by
    the Gmail agent/endpoints to pick the per-login-email mailbox (GMAIL_REFRESH_TOKENS).

    The token is still fully verified (invalid/expired tokens 401); a valid token that
    simply lacks an email claim returns "" so callers unrelated to Gmail aren't blocked.
    The Gmail path treats "" as "no mailbox" and reports it clearly.
    """
    payload = _decode_verified(credentials)
    return (payload.get("email") or "").strip().lower()


def _profile_status(user_id: str) -> str:
    now = time.monotonic()
    cached = _status_cache.get(user_id)
    if cached and now - cached[1] < _STATUS_TTL_SECONDS:
        return cached[0]
    res = _service_client().table("profiles").select("status").eq("id", user_id).execute()
    # Users created before the approval flow may lack a profiles row; they are
    # grandfathered in as approved.
    status = res.data[0].get("status") or "approved" if res.data else "approved"
    _status_cache[user_id] = (status, now)
    return status


def invalidate_status_cache(user_id: str) -> None:
    _status_cache.pop(user_id, None)


def get_user_source_access(user_id: str) -> str | None:
    """Return the raw profiles."Source access" string for a user (cached).

    This is the comma-separated list of data sources (e.g. "Sales, Inventory,
    Overdue, Order") the user is permitted to query via the SQL agent. Returns
    None when the profile has no value set — the caller treats that as "no
    data-source access". Failures are swallowed and cached as None so a
    transient DB error denies access rather than crashing the chat request.
    """
    now = time.monotonic()
    cached = _source_access_cache.get(user_id)
    if cached and now - cached[1] < _STATUS_TTL_SECONDS:
        return cached[0]
    try:
        # The column name contains a space, so it must be quoted for PostgREST.
        res = _service_client().table("profiles").select('"Source access"').eq("id", user_id).execute()
        value = res.data[0].get("Source access") if res.data else None
    except Exception:
        value = None
    _source_access_cache[user_id] = (value, now)
    return value


def invalidate_source_access_cache(user_id: str) -> None:
    _source_access_cache.pop(user_id, None)


def verify_approved_user(credentials: HTTPAuthorizationCredentials = Security(_security)) -> str:
    """verify_token + admin-approval gate: new signups start as 'pending' in
    public.profiles and get 403 here until an admin approves them."""
    user_id = verify_token(credentials)
    status = _profile_status(user_id)
    if status != "approved":
        detail = "Account rejected" if status == "rejected" else "Account pending admin approval"
        raise HTTPException(status_code=403, detail=detail)
    return user_id