src/app_helpers/auth.py

"""OAuth / whitelist helpers used by the Gradio submission flow."""

from __future__ import annotations

import logging

import gradio as gr
from huggingface_hub import whoami
from huggingface_hub.errors import HfHubHTTPError

from src.envs import SIZE_WHITELIST, WHITELIST

logger = logging.getLogger(__name__)


def check_intel_org(oauth_token: gr.OAuthToken | None) -> tuple[bool, bool]:
    """Return ``(is_authenticated, is_intel)``.

    Raises :class:`HfHubHTTPError` if the token is expired/invalid.
    """
    if oauth_token is None:
        return False, False
    info = whoami(token=oauth_token.token)
    username = info.get("name", "unknown")
    orgs = [org["name"] for org in info.get("orgs", [])]
    logger.info("[Auth] user=%s  orgs=%s", username, orgs)
    return True, "intel" in [o.lower() for o in orgs]


def is_whitelisted(oauth_token: gr.OAuthToken | None) -> bool:
    """Whether the authenticated user is in the submission whitelist."""
    if oauth_token is None or not oauth_token.token:
        return False
    if WHITELIST is None:
        # No whitelist configured → everyone with a valid token qualifies
        return True
    try:
        info = whoami(token=oauth_token.token)
        return info.get("name", "unknown") in WHITELIST
    except Exception:
        return False


def is_allowed_to_resubmit_failed(username: str) -> bool:
    """Whether *username* may re-submit a model that previously Failed.

    Mirrors the :data:`SIZE_WHITELIST` semantics:

    * ``SIZE_WHITELIST is None`` → everyone is allowed.
    * Otherwise only users in the set are allowed.
    """
    if SIZE_WHITELIST is None:
        return True
    return bool(username) and username in SIZE_WHITELIST


def process_user_on_load(oauth_token: gr.OAuthToken | None):
    """Page-load callback: log user info and show the hardware selector for
    ``SIZE_WHITELIST`` users.
    """
    show_hw = False
    username = ""
    if oauth_token is not None and oauth_token.token:
        try:
            info = whoami(token=oauth_token.token)
            username = info.get("name", "")
            orgs = [org["name"] for org in info.get("orgs", [])]
            logger.info("[Auth][PageLoad] user=%s  orgs=%s", username, orgs)
            if SIZE_WHITELIST is not None:
                show_hw = username in SIZE_WHITELIST
        except Exception as e:
            logger.warning("[Auth][PageLoad] whoami failed: %s", e)

    if username:
        user_html = (
            f'<span style="font-size:0.85rem;margin-left:4px;color:#374151;">'
            f'👤 <strong>{username}</strong></span>'
        )
        user_html_quant = user_html
        user_html_eval = user_html
    else:
        user_html_quant = (
            '<span style="color:#9ca3af;font-size:0.85rem;margin-left:4px;">'
            'Unauthorized — <a href="#" class="auth-trigger-link" data-tab="quant" '
            'style="color:#3b82f6;cursor:pointer;text-decoration:underline;">'
            'Click to authorize</a></span>'
        )
        user_html_eval = (
            '<span style="color:#9ca3af;font-size:0.85rem;margin-left:4px;">'
            'Unauthorized — <a href="#" class="auth-trigger-link" data-tab="eval" '
            'style="color:#3b82f6;cursor:pointer;text-decoration:underline;">'
            'Click to authorize</a></span>'
        )
    return (
        gr.update(visible=show_hw),           # hw_col_quant
        gr.update(visible=show_hw),           # hw_col_eval
        username,                             # username_state
        gr.update(value=user_html_quant),     # current_user_quant
        gr.update(value=user_html_eval),      # current_user_eval
    )


def check_auth_for_submit(oauth_token: gr.OAuthToken | None) -> str:
    """Pre-check auth before submission. Returns ``"TRIGGER"`` or ``"OK"``."""
    if oauth_token is None or not oauth_token.token:
        return "TRIGGER"
    try:
        whoami(token=oauth_token.token)
        return "OK"
    except (HfHubHTTPError, Exception):
        logger.warning("[Auth] token check failed (likely expired)", exc_info=True)
        return "TRIGGER"


def sync_username(oauth_token: gr.OAuthToken | None = None):
    """Refresh ``username_state`` after a successful auth.

    Returns :func:`gr.update` (no change) if the token is unavailable or the
    call fails.
    """
    if oauth_token is None or not oauth_token.token:
        return gr.update()
    try:
        info = whoami(token=oauth_token.token)
        return info.get("name", "")
    except Exception:
        return gr.update()