packages/core/src/utils/auth.py

# ============================================================================
# agent verf - Authentication & User Management
# Version: 0.1.0
# Last Updated: 2026-01-22
#
# Handles user management for dual-mode operation.
# Users authenticate via Google OAuth (Supabase Auth).
#
# User Tiers:
#   - Regular users: Free mode only
#   - Premium users: Venice mode by default, can toggle to Free
# ============================================================================

import os
import structlog
from typing import Optional
from pydantic import BaseModel
from enum import Enum
from supabase import AsyncClient, create_async_client

from src.utils.llm import VerificationMode

logger = structlog.get_logger()


# ============================================================================
# User Models
# ============================================================================

class User(BaseModel):
    """User account data."""
    id: str
    google_id: str
    email: Optional[str] = None
    display_name: Optional[str] = None
    avatar_url: Optional[str] = None
    is_premium: bool = False
    preferred_mode: str = "free"

    def get_verification_mode(self) -> VerificationMode:
        """Get the user's verification mode."""
        if not self.is_premium:
            return VerificationMode.FREE
        return VerificationMode(self.preferred_mode)


class UserCreate(BaseModel):
    """Data for creating a new user."""
    google_id: str
    email: Optional[str] = None
    display_name: Optional[str] = None
    avatar_url: Optional[str] = None


class UserUpdate(BaseModel):
    """Data for updating a user."""
    preferred_mode: Optional[str] = None
    display_name: Optional[str] = None
    avatar_url: Optional[str] = None


# ============================================================================
# User Service
# ============================================================================

class UserService:
    """
    Service for managing users in Supabase.

    Handles:
      - User creation on first login
      - User lookups by Google ID
      - Premium mode toggling
    """

    def __init__(self):
        """Initialize with Supabase client."""
        self._client: Optional[AsyncClient] = None

    async def _get_client(self) -> AsyncClient:
        """Get or create the Supabase async client."""
        if self._client is None:
            url = os.getenv("SUPABASE_URL", "")
            key = os.getenv("SUPABASE_SERVICE_KEY", "")

            if not url.startswith("http"):
                url = f"https://{url}"

            self._client = await create_async_client(url, key)

        return self._client

    async def get_or_create_user(
        self,
        google_id: str,
        email: Optional[str] = None,
        display_name: Optional[str] = None,
        avatar_url: Optional[str] = None,
    ) -> User:
        """
        Get existing user or create new one.

        Called after successful Google OAuth login.
        """
        client = await self._get_client()

        # Try to find existing user
        result = await client.table("users").select("*").eq("google_id", google_id).execute()

        if result.data:
            user_data = result.data[0]
            logger.info("User found", user_id=user_data["id"], google_id=google_id)
            return User(**user_data)

        # Create new user
        new_user = {
            "google_id": google_id,
            "email": email,
            "display_name": display_name,
            "avatar_url": avatar_url,
            "is_premium": False,
            "preferred_mode": "free",
        }

        result = await client.table("users").insert(new_user).execute()
        user_data = result.data[0]

        logger.info("User created", user_id=user_data["id"], google_id=google_id)
        return User(**user_data)

    async def get_user_by_id(self, user_id: str) -> Optional[User]:
        """Get user by their UUID."""
        client = await self._get_client()

        result = await client.table("users").select("*").eq("id", user_id).execute()

        if result.data:
            return User(**result.data[0])
        return None

    async def get_user_by_google_id(self, google_id: str) -> Optional[User]:
        """Get user by their Google ID."""
        client = await self._get_client()

        result = await client.table("users").select("*").eq("google_id", google_id).execute()

        if result.data:
            return User(**result.data[0])
        return None

    async def update_user_mode(self, user_id: str, mode: VerificationMode) -> Optional[User]:
        """
        Update user's preferred mode.

        Only premium users can change their mode.
        """
        client = await self._get_client()

        # First check if user is premium
        user = await self.get_user_by_id(user_id)
        if not user:
            logger.warning("User not found", user_id=user_id)
            return None

        if not user.is_premium:
            logger.warning("Non-premium user attempted mode change", user_id=user_id)
            return None

        # Update the mode
        result = await client.table("users").update({
            "preferred_mode": mode.value
        }).eq("id", user_id).execute()

        if result.data:
            logger.info("User mode updated", user_id=user_id, mode=mode.value)
            return User(**result.data[0])

        return None

    async def set_premium_status(self, user_id: str, is_premium: bool) -> Optional[User]:
        """
        Set user's premium status.

        Admin-only operation.
        """
        client = await self._get_client()

        update_data = {"is_premium": is_premium}

        # If removing premium, reset to free mode
        if not is_premium:
            update_data["preferred_mode"] = "free"

        result = await client.table("users").update(update_data).eq("id", user_id).execute()

        if result.data:
            logger.info("User premium status updated", user_id=user_id, is_premium=is_premium)
            return User(**result.data[0])

        return None

    async def get_verification_mode(self, user_id: Optional[str]) -> VerificationMode:
        """
        Get the verification mode for a user.

        Returns FREE if user is not found or not premium.
        """
        if not user_id:
            return VerificationMode.FREE

        user = await self.get_user_by_id(user_id)
        if not user:
            return VerificationMode.FREE

        return user.get_verification_mode()


# ============================================================================
# Singleton Instance
# ============================================================================

_user_service: Optional[UserService] = None


def get_user_service() -> UserService:
    """Get the singleton UserService instance."""
    global _user_service
    if _user_service is None:
        _user_service = UserService()
    return _user_service