WealthWise/backend/app/schemas/token.py

"""Pydantic models for JWT token handling and user context.

This module defines the data structures for:
- TokenPayload: Full JWT payload from Supabase
- UserContext: Clean user representation for API context
"""

from typing import Any, Optional

from pydantic import BaseModel, Field


class TokenPayload(BaseModel):
    """Supabase JWT token payload structure.
    
    This model represents the complete JWT payload that Supabase Auth
    includes in the access token. It contains all claims including
    standard JWT claims and Supabase-specific claims.
    
    Attributes:
        sub: Subject (user UUID)
        exp: Expiration timestamp (Unix)
        iat: Issued at timestamp (Unix)
        aud: Audience (should be "authenticated")
        email: User's email address
        phone: User's phone number (if available)
        app_metadata: Application-specific metadata (includes role)
        user_metadata: User-specific metadata
        role: User's database role (e.g., "authenticated", "anon")
        aal: Authenticator assurance level (aal1, aal2)
        amr: Authentication methods reference
        session_id: Session UUID
        is_anonymous: Whether this is an anonymous user
    """
    
    # Standard JWT claims
    sub: str = Field(description="Subject (user UUID)")
    exp: Optional[int] = Field(default=None, description="Expiration timestamp (Unix)")
    iat: Optional[int] = Field(default=None, description="Issued at timestamp (Unix)")
    iss: Optional[str] = Field(default=None, description="Issuer")
    aud: str = Field(description="Audience (should be 'authenticated')")
    
    # User identity claims
    email: Optional[str] = Field(default=None, description="User's email address")
    phone: Optional[str] = Field(default=None, description="User's phone number")
    email_confirmed_at: Optional[str] = Field(default=None, description="Email confirmation timestamp")
    phone_confirmed_at: Optional[str] = Field(default=None, description="Phone confirmation timestamp")
    
    # Supabase metadata
    app_metadata: dict[str, Any] = Field(
        default_factory=dict,
        description="Application-specific metadata (provider, role, etc.)"
    )
    user_metadata: dict[str, Any] = Field(
        default_factory=dict,
        description="User-specific metadata"
    )
    
    # Supabase-specific claims
    role: Optional[str] = Field(default=None, description="Database role")
    aal: Optional[str] = Field(default=None, description="Authenticator assurance level")
    amr: Optional[list[dict[str, Any]]] = Field(default=None, description="Authentication methods reference")
    session_id: Optional[str] = Field(default=None, description="Session UUID")
    is_anonymous: Optional[bool] = Field(default=False, description="Whether user is anonymous")
    
    class Config:
        """Pydantic configuration."""
        extra = "allow"  # Allow additional claims not explicitly defined


class UserContext(BaseModel):
    """Clean user context for API endpoints.
    
    This model represents the authenticated user information that
    is passed to API endpoints via the get_current_user dependency.
    It contains only the essential information needed by most endpoints.
    
    Attributes:
        id: User UUID from the 'sub' claim
        email: User's email address
        role: User's role (extracted from app_metadata or role claim)
        is_anonymous: Whether this is an anonymous user
        session_id: Current session ID
    """
    
    id: str = Field(description="User UUID")
    email: Optional[str] = Field(default=None, description="User's email address")
    role: str = Field(default="authenticated", description="User's role")
    is_anonymous: bool = Field(default=False, description="Whether user is anonymous")
    session_id: Optional[str] = Field(default=None, description="Current session ID")
    
    @classmethod
    def from_token_payload(cls, payload: TokenPayload) -> "UserContext":
        """Create UserContext from TokenPayload.
        
        Extracts user information from the full JWT payload and creates
        a clean UserContext object. Role is extracted from app_metadata
        if available, otherwise defaults to the role claim or 'authenticated'.
        
        Args:
            payload: Validated TokenPayload from JWT
            
        Returns:
            UserContext with extracted user information
        """
        # Extract role from app_metadata or fall back to role claim
        role = payload.app_metadata.get("role", payload.role or "authenticated")
        
        return cls(
            id=payload.sub,
            email=payload.email,
            role=role,
            is_anonymous=payload.is_anonymous or False,
            session_id=payload.session_id,
        )
    
    class Config:
        """Pydantic configuration."""
        json_schema_extra = {
            "example": {
                "id": "550e8400-e29b-41d4-a716-446655440000",
                "email": "user@wealthwise.app",
                "role": "authenticated",
                "is_anonymous": False,
                "session_id": "session-uuid-here",
            }
        }