"""Multi-user tenant authentication with Google OAuth, CSRF protection, and automatic tenant provisioning."""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_utils_auth.ipynb.

# %% auto #0
__all__ = ['logger', 'ROLE_HIERARCHY', 'DEFAULT_SKIP_AUTH', 'SessionConfig', 'SlidingSessionMiddleware',
           'create_session_middleware', 'has_min_role', 'get_user_role', 'require_role', 'invalidate_auth_cache',
           'create_auth_beforeware', 'get_google_oauth_client', 'generate_oauth_state', 'verify_oauth_state',
           'create_or_get_global_user', 'get_user_membership', 'verify_membership', 'provision_new_user',
           'create_user_session', 'get_current_user', 'clear_session', 'auth_redirect', 'route_user_after_login',
           'require_tenant_access', 'handle_login_request', 'handle_oauth_callback', 'handle_logout']

# %% ../nbs/04_utils_auth.ipynb #99e1cd05
from nbdev.showdoc import show_doc
from fastsql import *
from fastcore.utils import *
from fh_saas.db_host import (
    timestamp, gen_id, 
    GlobalUser, TenantCatalog, Membership, HostAuditLog,
    HostDatabase
)
from fh_saas.db_tenant import (
    get_or_create_tenant_db, 
    init_tenant_core_schema, 
    TenantUser
)
from fasthtml.oauth import GoogleAppClient, redir_url
from starlette.responses import RedirectResponse, Response
import os
import uuid
import json
import logging
import time
from dotenv import load_dotenv
from dataclasses import dataclass
from starlette.middleware.sessions import SessionMiddleware
from starlette.datastructures import MutableHeaders
from itsdangerous import URLSafeTimedSerializer, BadSignature
import base64

logger = logging.getLogger(__name__)

load_dotenv()# Module-level logger - configured by app via configure_logging()


# %% ../nbs/04_utils_auth.ipynb #e215776c
@dataclass
class SessionConfig:
    """Configuration for sliding session behavior.
    
    Attributes:
        max_age: Session expires after this many seconds of inactivity. Default: 3600 (1 hour)
        sliding: If True, refresh session expiry on each request. Default: True
        absolute_max: Optional hard limit in seconds regardless of activity. Default: None
        secure: If True, cookie only sent over HTTPS. Default: True
        same_site: SameSite cookie policy ('lax', 'strict', 'none'). Default: 'lax'
        http_only: If True, cookie not accessible via JavaScript. Default: True
        
    Example:
        >>> config = SessionConfig()  # 1 hour inactivity timeout, sliding enabled
        >>> config = SessionConfig(max_age=1800)  # 30 min inactivity timeout
        >>> config = SessionConfig(max_age=3600, absolute_max=86400)  # 1h sliding, 24h hard limit
    """
    max_age: int = 3600  # 1 hour inactivity timeout
    sliding: bool = True  # Refresh on each request
    absolute_max: int = None  # Optional hard limit (e.g., 86400 for 24h)
    secure: bool = True  # HTTPS-only in production
    same_site: str = 'lax'  # CSRF protection
    http_only: bool = True  # No JS access
    
    @classmethod
    def default(cls) -> 'SessionConfig':
        """Standard sliding session: 1 hour inactivity timeout."""
        return cls()
    
    @classmethod
    def strict(cls) -> 'SessionConfig':
        """Strict session: 30 min inactivity, 8 hour hard limit."""
        return cls(max_age=1800, absolute_max=28800)
    
    @classmethod
    def relaxed(cls) -> 'SessionConfig':
        """Relaxed session: 8 hour inactivity, 7 day hard limit."""
        return cls(max_age=28800, absolute_max=604800)

# %% ../nbs/04_utils_auth.ipynb #edc959ee
class SlidingSessionMiddleware(SessionMiddleware):
    """Session middleware with sliding expiry.
    
    Extends Starlette's SessionMiddleware to refresh the session cookie
    max_age on each request, implementing sliding session expiry.
    
    Sessions expire after `max_age` seconds of INACTIVITY, not from login time.
    
    Args:
        app: ASGI application
        secret_key: Secret key for signing cookies
        session_config: SessionConfig instance for timeout settings
        session_cookie: Cookie name. Default: 'session'
        path: Cookie path. Default: '/'
        
    Example:
        >>> from starlette.applications import Starlette
        >>> app = Starlette()
        >>> config = SessionConfig(max_age=3600)  # 1 hour inactivity
        >>> app = SlidingSessionMiddleware(app, secret_key='...', session_config=config)
    """
    
    def __init__(
        self,
        app,
        secret_key: str,
        session_config: SessionConfig = None,
        session_cookie: str = "session",
        path: str = "/",
    ):
        self.config = session_config or SessionConfig.default()
        
        # Initialize parent with config values
        super().__init__(
            app,
            secret_key=secret_key,
            session_cookie=session_cookie,
            max_age=self.config.max_age,
            path=path,
            same_site=self.config.same_site,
            https_only=self.config.secure,
        )
        self._config = self.config  # Store for access in __call__
    
    async def __call__(self, scope, receive, send):
        """Process request and refresh session cookie on response."""
        if scope["type"] not in ("http", "websocket"):
            await self.app(scope, receive, send)
            return
        
        # Load existing session
        connection = (scope["type"] == "http") and scope or None
        if connection:
            # Standard session loading from parent
            initial_session_was_empty = True
            scope["session"] = {}
            
            # Try to load existing session from cookie
            if self.session_cookie in scope.get("headers", []):
                pass  # Will be loaded by parent
        
        # Let parent handle the request
        async def send_wrapper(message):
            """Wrap send to refresh session cookie max_age."""
            if message["type"] == "http.response.start":
                # Check if session has user_id (authenticated)
                session = scope.get("session", {})
                
                if session.get("user_id") and self._config.sliding:
                    # Check absolute max if configured
                    if self._config.absolute_max:
                        session_started = session.get("session_started_at", 0)
                        if session_started and (time.time() - session_started > self._config.absolute_max):
                            # Absolute max exceeded - clear session
                            session.clear()
                            logger.info("Session absolute max exceeded, forcing logout")
                    
                    # Session is valid - cookie will be refreshed with new max_age
                    # The parent class handles setting the cookie with max_age
                    logger.debug(f"Sliding session refresh for user {session.get('user_id')}")
            
            await send(message)
        
        await super().__call__(scope, receive, send_wrapper)


def create_session_middleware(
    secret_key: str,
    session_config: SessionConfig = None,
    session_cookie: str = "session",
) -> SlidingSessionMiddleware:
    """Factory to create SlidingSessionMiddleware for FastHTML apps.
    
    Args:
        secret_key: Secret key for signing session cookies (required)
        session_config: SessionConfig instance. Default: SessionConfig.default()
        session_cookie: Cookie name. Default: 'session'
        
    Returns:
        Configured SlidingSessionMiddleware instance
        
    Example:
        >>> from fasthtml.common import FastHTML
        >>> 
        >>> # Create app WITHOUT default session middleware
        >>> app = FastHTML(sess_cls=None)  # Disable default sessions
        >>> 
        >>> # Add sliding session middleware
        >>> config = SessionConfig(max_age=3600)  # 1 hour inactivity
        >>> app = create_session_middleware('your-secret-key', config)(app)
        >>>
        >>> # Or wrap during app creation (recommended)
        >>> middleware = create_session_middleware('your-secret-key', config)
        >>> # Then configure FastHTML to use it
        
    Note:
        FastHTML's default SessionMiddleware needs to be disabled first.
        See integration documentation for patterns.
    """
    config = session_config or SessionConfig.default()
    
    def middleware_factory(app):
        return SlidingSessionMiddleware(
            app,
            secret_key=secret_key,
            session_config=config,
            session_cookie=session_cookie,
        )
    
    return middleware_factory

# %% ../nbs/04_utils_auth.ipynb #084765e5
# Role hierarchy: higher number = more permissions
ROLE_HIERARCHY = {
    'admin': 3,
    'editor': 2,
    'viewer': 1
}

def has_min_role(user: dict, required_role: str) -> bool:
    """Check if user meets the minimum role requirement.
    
    Args:
        user: User dict with 'role' field (from request.state.user)
        required_role: Minimum role needed ('admin', 'editor', 'viewer')
    
    Returns:
        True if user's role >= required_role in hierarchy
        
    Example:
        >>> user = {'role': 'editor'}
        >>> has_min_role(user, 'viewer')  # True - editor > viewer
        >>> has_min_role(user, 'admin')   # False - editor < admin
    """
    user_role = user.get('role', 'viewer')
    user_level = ROLE_HIERARCHY.get(user_role, 0)
    required_level = ROLE_HIERARCHY.get(required_role, 0)
    return user_level >= required_level


def get_user_role(session: dict, tenant_db: 'Database' = None) -> str:
    """Derive effective role from session and tenant database.
    
    Rules:
    1. Tenant owner (session['tenant_role'] == 'owner') → 'admin'
    2. System admin → 'admin'
    3. Otherwise → lookup TenantUser.local_role from tenant DB
    4. Fallback → None (user must be explicitly assigned a role)
    
    Args:
        session: User session dict
        tenant_db: Tenant database connection (optional)
    
    Returns:
        Effective role string: 'admin', 'editor', 'viewer', or None
    """
    # Owner is always admin
    if session.get('tenant_role') == 'owner':
        return 'admin'
    
    # System admin is always admin
    if session.get('is_sys_admin'):
        return 'admin'
    
    # Look up local_role from TenantUser
    if tenant_db:
        user_id = session.get('user_id')
        if user_id:
            try:
                tenant_db.conn.rollback()
                tenant_users = tenant_db.t.core_tenant_users
                all_users = tenant_users()
                matching = [u for u in all_users if u.id == user_id]
                if matching:
                    return matching[0].local_role
            except Exception as e:
                logger.warning(f"Failed to lookup TenantUser role: {e}")
    
    # No role found - user must be assigned explicitly
    return None

# %% ../nbs/04_utils_auth.ipynb #ca53191d
from functools import wraps
from starlette.responses import Response

def require_role(min_role: str):
    """Decorator to protect routes with minimum role requirement.
    
    Args:
        min_role: Minimum role required ('admin', 'editor', 'viewer')
    
    Returns:
        Decorator that checks request.state.user['role']
        Returns 403 Forbidden if user lacks required role
        
    Example:
        >>> @app.get('/admin/settings')
        >>> @require_role('admin')
        >>> def admin_settings(request):
        ...     return "Admin only content"
        
        >>> @app.get('/reports')  
        >>> @require_role('viewer')  # All authenticated users
        >>> def view_reports(request):
        ...     return "Reports"
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(request, *args, **kwargs):
            user = getattr(request.state, 'user', None)
            
            if not user:
                return Response(
                    content="Authentication required",
                    status_code=401
                )
            
            if not user.get('role'):
                return Response(
                    content="No role assigned. Contact your administrator.",
                    status_code=403
                )
            
            if not has_min_role(user, min_role):
                return Response(
                    content=f"Access denied. Required role: {min_role}",
                    status_code=403
                )
            
            # Handle both sync and async route functions
            result = func(request, *args, **kwargs)
            if hasattr(result, '__await__'):
                return await result
            return result
        
        return wrapper
    return decorator

# %% ../nbs/04_utils_auth.ipynb #3acff0ac
# Session cache key
_AUTH_CACHE_KEY = '_auth_cache'

def invalidate_auth_cache(session: dict):
    """Clear the auth cache from session.
    
    Call this when:
    - User role or permissions change
    - User is added/removed from tenant
    - Admin changes user's local_role
    
    Args:
        session: User session dict
        
    Example:
        >>> # After admin changes user role
        >>> tenant_user.local_role = 'editor'
        >>> tables['tenant_users'].update(tenant_user)
        >>> invalidate_auth_cache(session)  # Force fresh lookup
    """
    session.pop(_AUTH_CACHE_KEY, None)
    logger.debug("Auth cache invalidated")


def _get_cached_auth(session: dict, cache_ttl: int) -> dict | None:
    """Get cached auth data if valid and not expired.
    
    Args:
        session: User session dict
        cache_ttl: Cache time-to-live in seconds
        
    Returns:
        Cached auth dict with 'user' and 'tenant_id', or None if expired/missing
    """
    cache = session.get(_AUTH_CACHE_KEY)
    if not cache:
        return None
    
    cached_at = cache.get('cached_at', 0)
    if time.time() - cached_at > cache_ttl:
        logger.debug("Auth cache expired")
        return None
    
    return cache


def _set_auth_cache(session: dict, user: dict, tenant_id: str):
    """Store auth data in session cache.
    
    Args:
        session: User session dict
        user: User dict to cache
        tenant_id: Tenant ID to cache
    """
    session[_AUTH_CACHE_KEY] = {
        'user': user.copy(),
        'tenant_id': tenant_id,
        'cached_at': time.time()
    }
    logger.debug(f"Auth cache set for user {user.get('user_id')}")

# %% ../nbs/04_utils_auth.ipynb #17a4dd9c
from fasthtml.common import Beforeware
from typing import Callable, Any

# Default URL patterns to skip authentication
DEFAULT_SKIP_AUTH = [
    r'/login.*',
    r'/logout',
    r'/oauth/callback',
    r'/auth/callback',
    r'/health',
    r'/favicon\.ico',
    r'/static/.*',
]

def create_auth_beforeware(
    redirect_path: str = '/login',
    session_key: str = 'user_id',
    skip: list[str] = None,
    include_defaults: bool = True,
    setup_tenant_db: bool = True,
    schema_init: Callable[[Database], dict[str, Any]] = None,
    session_cache: bool = False,
    session_cache_ttl: int = 300,
    require_subscription: bool = False,
    subscription_redirect: str = None,
    grace_period_days: int = 3,
    session_config: SessionConfig = None,
):
    """Create Beforeware that checks for authenticated session and sets up request.state.
    
    Args:
        redirect_path: Where to redirect unauthenticated users
        session_key: Session key for user ID
        skip: List of regex patterns to skip auth
        include_defaults: Include default skip patterns
        setup_tenant_db: Auto-setup tenant database on request.state
        schema_init: Optional callback to initialize tables dict.
                     Signature: (tenant_db: Database) -> dict[str, Table]
                     Result stored in request.state.tables
        session_cache: Enable caching user dict in session to reduce DB queries.
                       Recommended for HTMX-heavy apps. Default: False
        session_cache_ttl: Cache TTL in seconds. Default: 300 (5 minutes)
        require_subscription: If True, check for active subscription and return 402
                             if not found. Default: False
        subscription_redirect: Optional URL to redirect if subscription required.
                              If None, returns 402 Payment Required response.
        grace_period_days: Days to allow access after payment failure (default: 3)
        session_config: Optional SessionConfig for absolute timeout enforcement.
                       Note: Sliding expiry requires SlidingSessionMiddleware.
                       This parameter only enforces absolute_max if configured.
    
    Returns:
        Beforeware instance for FastHTML apps
        
    Sets on request.state:
        - user: dict with user_id, email, tenant_id, role, is_owner
        - tenant_id: str
        - tenant_db: Database connection
        - tables: dict of Table objects (if schema_init provided)
        - subscription: Subscription object (if require_subscription enabled)
        
    Example:
        >>> # Basic usage
        >>> beforeware = create_auth_beforeware()
        
        >>> # With session caching for HTMX apps
        >>> beforeware = create_auth_beforeware(
        ...     session_cache=True,
        ...     session_cache_ttl=300
        ... )
        
        >>> # With schema initialization
        >>> def get_app_tables(db):
        ...     return {'users': db.create(User, pk='id')}
        >>> beforeware = create_auth_beforeware(schema_init=get_app_tables)
        
        >>> # With subscription requirement
        >>> beforeware = create_auth_beforeware(
        ...     require_subscription=True,
        ...     subscription_redirect='/pricing'
        ... )
    """
    skip_patterns = []
    if include_defaults:
        skip_patterns.extend(DEFAULT_SKIP_AUTH)
    if skip:
        skip_patterns.extend(skip)
    
    def check_auth(req, sess):
        # Check absolute session timeout if session_config provided
        if session_config and session_config.absolute_max:
            session_started = sess.get('session_started_at', 0)
            if session_started and (time.time() - session_started > session_config.absolute_max):
                logger.info(f"Session absolute max ({session_config.absolute_max}s) exceeded, forcing logout")
                clear_session(sess)
                return auth_redirect(req, redirect_path)
        
        if session_key not in sess:
            return auth_redirect(req, redirect_path)
        
        if setup_tenant_db:
            # Check session cache first (if enabled)
            cache_hit = False
            if session_cache:
                cached = _get_cached_auth(sess, session_cache_ttl)
                if cached:
                    cache_hit = True
                    user = cached['user']
                    req.state.user = user
                    req.state.tenant_id = cached['tenant_id']
                    
                    # Still need to create tenant_db connection (lightweight)
                    if user.get('tenant_id') and not user.get('is_sys_admin'):
                        try:
                            req.state.tenant_db = get_or_create_tenant_db(user['tenant_id'])
                        except Exception as e:
                            logger.error(f"Failed to setup tenant_db from cache: {e}")
                            req.state.tenant_db = None
                    else:
                        req.state.tenant_db = None
            
            # Cache miss - do full DB lookup
            if not cache_hit:
                user = get_current_user(sess)
                if user:
                    req.state.tenant_id = user.get('tenant_id')
                    req.state.tenant_db = None
                    
                    if user.get('tenant_id') and not user.get('is_sys_admin'):
                        try:
                            host_db = HostDatabase.from_env()
                            if verify_membership(host_db, user['user_id'], user['tenant_id']):
                                req.state.tenant_db = get_or_create_tenant_db(user['tenant_id'])
                            else:
                                logger.warning(f"Invalid membership for user {user['user_id']}")
                        except Exception as e:
                            logger.error(f"Failed to setup tenant_db: {e}")
                    
                    # Derive effective role from session + tenant DB
                    role = get_user_role(sess, req.state.tenant_db)
                    user['role'] = role
                    user['is_owner'] = sess.get('tenant_role') == 'owner'
                    req.state.user = user
                    
                    # Update session cache (if enabled)
                    if session_cache:
                        _set_auth_cache(sess, user, user.get('tenant_id'))
            
            # Auto-initialize tables if schema_init provided
            if schema_init and req.state.tenant_db:
                try:
                    req.state.tables = schema_init(req.state.tenant_db)
                except Exception as e:
                    logger.error(f"Failed to initialize schema: {e}")
                    req.state.tables = {}
            else:
                req.state.tables = {}
            
            # Check subscription requirement
            if require_subscription and hasattr(req.state, 'tenant_id') and req.state.tenant_id:
                try:
                    from fh_saas.utils_stripe import get_active_subscription
                    subscription = get_active_subscription(
                        req.state.tenant_id,
                        grace_period_days=grace_period_days
                    )
                    req.state.subscription = subscription
                    
                    if not subscription:
                        logger.warning(f"Subscription required for tenant {req.state.tenant_id}")
                        if subscription_redirect:
                            return auth_redirect(req, subscription_redirect)
                        return Response(
                            content='Payment required',
                            status_code=402,
                            media_type='text/plain'
                        )
                except ImportError:
                    logger.error("utils_stripe not available for subscription check")
                    req.state.subscription = None
                except Exception as e:
                    logger.error(f"Error checking subscription: {e}")
                    req.state.subscription = None
    
    return Beforeware(check_auth, skip=skip_patterns)

# %% ../nbs/04_utils_auth.ipynb #0d965490
def get_google_oauth_client():
    """Initialize Google OAuth client with credentials from environment."""
    client_id = os.getenv('GOOGLE_CLIENT_ID')
    client_secret = os.getenv('GOOGLE_CLIENT_SECRET')
    
    if not client_id or not client_secret:
        raise ValueError("Missing GOOGLE_CLIENT_ID or GOOGLE_CLIENT_SECRET in .env")
    
    return GoogleAppClient(client_id=client_id, client_secret=client_secret)

# %% ../nbs/04_utils_auth.ipynb #951b6922
def generate_oauth_state():
    """Generate cryptographically secure random state token for CSRF protection."""
    return uuid.uuid4().hex


def verify_oauth_state(session: dict, callback_state: str):
    """Verify OAuth callback state matches stored session state (CSRF protection)."""
    stored_state = session.get('oauth_state')
    
    if not stored_state:
        raise ValueError("CSRF validation failed: No state in session")
    
    if stored_state != callback_state:
        raise ValueError("CSRF validation failed: State mismatch")
    
    session.pop('oauth_state', None)  # One-time use

# %% ../nbs/04_utils_auth.ipynb #9eb258ad
def create_or_get_global_user(host_db: HostDatabase, oauth_id: str, email: str, oauth_info: dict = None):
    """Create or retrieve GlobalUser from host database."""
    try:
        host_db.rollback()
        all_users = host_db.global_users()
        existing = [u for u in all_users if u.oauth_id == oauth_id]
        
        if existing:
            user = existing[0]
            user.last_login = timestamp()
            host_db.global_users.update(user)
            logger.info(f'User login: {email}', extra={'user_id': user.id, 'email': email})
            return user
        
        new_user = GlobalUser(
            id=gen_id(),
            email=email,
            oauth_id=oauth_id,
            created_at=timestamp(),
            last_login=timestamp()
        )
        host_db.global_users.insert(new_user)
        logger.info(f'New user created: {email}', extra={'user_id': new_user.id, 'email': email})
        return new_user
        
    except Exception as e:
        host_db.rollback()
        logger.error(f'Failed to create/get user {email}: {e}', exc_info=True)
        raise


def get_user_membership(host_db: HostDatabase, user_id: str):
    """Get single active membership for user."""
    host_db.rollback()
    all_memberships = host_db.memberships()
    active = [m for m in all_memberships if m.user_id == user_id and m.is_active]
    return active[0] if active else None


def verify_membership(host_db: HostDatabase, user_id: str, tenant_id: str) -> bool:
    """Verify user has active membership for specific tenant."""
    host_db.rollback()
    all_memberships = host_db.memberships()
    valid = [m for m in all_memberships if m.user_id == user_id and m.tenant_id == tenant_id and m.is_active]
    return len(valid) > 0

# %% ../nbs/04_utils_auth.ipynb #7d74645e
def provision_new_user(host_db: HostDatabase, global_user: GlobalUser) -> str:
    """Auto-provision new tenant for first-time user."""
    tenant_id = gen_id()
    username = global_user.email.split('@')[0]
    tenant_name = f"{username}'s Workspace"
    tenant_db = None
    
    try:
        logger.info(f'Starting tenant provisioning for {global_user.email}',
                   extra={'tenant_id': tenant_id, 'user_id': global_user.id})
        
        # Create physical tenant database and register in catalog
        tenant_db = get_or_create_tenant_db(tenant_id, tenant_name)
        
        # Initialize core tenant schema
        core_tables = init_tenant_core_schema(tenant_db)
        
        # Create TenantUser profile
        tenant_user = TenantUser(
            id=global_user.id,
            display_name=username,
            local_role='admin',
            created_at=timestamp()
        )
        core_tables['tenant_users'].insert(tenant_user)
        tenant_db.conn.commit()  # Commit tenant changes
        
        # Create membership in host database
        membership = Membership(
            id=gen_id(),
            user_id=global_user.id,
            tenant_id=tenant_id,
            profile_id=global_user.id,
            role='owner',
            created_at=timestamp()
        )
        host_db.memberships.insert(membership)
        
        # Log provisioning event
        audit_log = HostAuditLog(
            id=gen_id(),
            actor_user_id=global_user.id,
            event_type='tenant_provisioned',
            target_id=tenant_id,
            details=json.dumps({'tenant_name': tenant_name, 'plan_tier': 'free', 'user_email': global_user.email}),
            created_at=timestamp()
        )
        host_db.audit_logs.insert(audit_log)
        
        host_db.commit()
        logger.info(f'Tenant provisioned: {tenant_name}',
                   extra={'tenant_id': tenant_id, 'tenant_name': tenant_name, 'user_id': global_user.id})
        return tenant_id
        
    except Exception as e:
        host_db.rollback()
        if tenant_db:
            try:
                tenant_db.conn.rollback()
            except Exception:
                pass
        logger.error(f'Tenant provisioning failed for {global_user.email}: {e}',
                    extra={'tenant_id': tenant_id, 'user_id': global_user.id}, exc_info=True)
        raise Exception(f"Failed to provision tenant for {global_user.email}: {str(e)}") from e
    finally:
        # Always close tenant_db connection to prevent connection leaks
        if tenant_db:
            try:
                tenant_db.conn.close()
                tenant_db.engine.dispose()
            except Exception:
                pass

# %% ../nbs/04_utils_auth.ipynb #5b1f2810
def create_user_session(session: dict, global_user: GlobalUser, membership: Membership):
    """Create authenticated session after successful OAuth login.
    
    Sets session keys for user identity and tracking:
    - user_id, email, tenant_id, tenant_role, is_sys_admin: Identity
    - login_at: Timestamp when user logged in
    - session_started_at: Unix timestamp for absolute session timeout tracking
    """
    session['user_id'] = global_user.id
    session['email'] = global_user.email
    session['tenant_id'] = membership.tenant_id
    session['tenant_role'] = membership.role
    session['is_sys_admin'] = global_user.is_sys_admin
    session['login_at'] = timestamp()
    session['session_started_at'] = time.time()  # For absolute max timeout tracking


def get_current_user(session: dict) -> dict | None:
    """Extract current user info from session.
    
    Returns:
        dict with keys: user_id, email, tenant_id, tenant_role, is_sys_admin
        
    Note:
        The 'role' and 'is_owner' fields are added by create_auth_beforeware
        after deriving the effective role from TenantUser.local_role.
        Access via request.state.user['role'] in routes.
    """
    if 'user_id' not in session:
        return None
    
    return {
        'user_id': session.get('user_id'),
        'email': session.get('email'),
        'tenant_id': session.get('tenant_id'),
        'tenant_role': session.get('tenant_role'),
        'is_sys_admin': session.get('is_sys_admin', False)
    }


def clear_session(session: dict):
    """Clear all session data (logout)."""
    session.clear()

# %% ../nbs/04_utils_auth.ipynb #30f53685
def auth_redirect(request, redirect_url: str = '/login'):
    """
    HTMX-aware redirect for authentication flows.
    
    When HTMX makes a partial request and receives a standard redirect (302/303),
    it follows the redirect and swaps the response into the target element.
    This causes the login page to appear inside the partial content area.
    
    This function detects HTMX requests and uses the `HX-Redirect` header
    to trigger a full page navigation instead.
    
    Args:
        request: Starlette request object
        redirect_url: URL to redirect to (default: '/login')
        
    Returns:
        Response with appropriate redirect mechanism
        
    Example:
        ```python
        @app.get('/dashboard')
        def dashboard(request):
            if not get_current_user(request.session):
                return auth_redirect(request)
            return render_dashboard()
        ```
    """
    # Check if this is an HTMX request
    if 'HX-Request' in request.headers:
        # Return 200 with HX-Redirect header - HTMX will do full page redirect
        response = Response(status_code=200)
        response.headers['HX-Redirect'] = redirect_url
        logger.debug(f"HTMX auth redirect to {redirect_url}")
        return response
    else:
        # Regular redirect for non-HTMX requests
        logger.debug(f"Standard auth redirect to {redirect_url}")
        return RedirectResponse(redirect_url, status_code=303)


def route_user_after_login(global_user: GlobalUser, membership: Membership = None) -> str:
    """Determine redirect URL based on user type and membership."""
    if global_user.is_sys_admin:
        return '/admin/dashboard'
    if membership:
        return '/dashboard'
    raise ValueError(f"User {global_user.email} has no membership")


def require_tenant_access(request_or_session):
    """Get tenant database with membership validation."""
    # Check if request.state.tenant_db already set by beforeware
    if hasattr(request_or_session, 'state'):
        if hasattr(request_or_session.state, 'tenant_db') and request_or_session.state.tenant_db:
            return request_or_session.state.tenant_db
        session = getattr(request_or_session, 'session', {})
    else:
        session = request_or_session
    
    user = get_current_user(session)
    if not user:
        raise ValueError("Authentication required")
    
    host_db = HostDatabase.from_env()
    
    if not verify_membership(host_db, user['user_id'], user['tenant_id']):
        raise PermissionError(f"Access denied for user {user['user_id']} to tenant {user['tenant_id']}")
    
    return get_or_create_tenant_db(user['tenant_id'])

# %% ../nbs/04_utils_auth.ipynb #a4b7b0cc
def handle_login_request(request, session):
    """Generate Google OAuth URL with CSRF state protection."""
    logger.debug('Login request initiated')
    
    # Generate CSRF state token
    state = generate_oauth_state()
    session['oauth_state'] = state
    
    # Get OAuth client and generate login link
    client = get_google_oauth_client()
    redirect_uri = redir_url(request, '/auth/callback')
    login_link = client.login_link(redirect_uri=redirect_uri, state=state)
    
    return login_link


def handle_oauth_callback(code: str, state: str, request, session):
    """Complete OAuth flow: CSRF verify → user info → provision → session → redirect."""
    logger.debug('OAuth callback received')
    
    # Step 1: CSRF validation (CRITICAL - must be first)
    verify_oauth_state(session, state)
    
    # Step 2: Exchange authorization code for user info
    client = get_google_oauth_client()
    redirect_uri = redir_url(request, '/auth/callback')
    user_info = client.retr_info(code, redirect_uri)
    
    # Step 3: Get host database instance (singleton - no need to close)
    host_db = HostDatabase.from_env()
    
    # Step 4: Create or get GlobalUser
    oauth_id = user_info[client.id_key]  # Google 'sub' field
    email = user_info.get('email', '')
    global_user = create_or_get_global_user(host_db, oauth_id, email, user_info)
    
    # Step 5: Check for existing membership
    membership = get_user_membership(host_db, global_user.id)
    
    # Step 6: Auto-provision if new user (no membership)
    if not membership and not global_user.is_sys_admin:
        tenant_id = provision_new_user(host_db, global_user)
        membership = get_user_membership(host_db, global_user.id)
    
    # Step 7: Create session (skip for sys admin - no tenant)
    if membership:
        create_user_session(session, global_user, membership)
    else:
        # System admin - minimal session
        session['user_id'] = global_user.id
        session['email'] = global_user.email
        session['is_sys_admin'] = True
        session['login_at'] = timestamp()
    
    # Step 8: Route to appropriate dashboard
    redirect_url = route_user_after_login(global_user, membership)
    logger.info(f'OAuth complete, redirecting to {redirect_url}', extra={'email': email})
    return RedirectResponse(redirect_url, status_code=303)


def handle_logout(session):
    """Clear session and redirect to login page."""
    email = session.get('email', 'unknown')
    clear_session(session)
    logger.info(f'User logged out: {email}')
    return RedirectResponse('/login', status_code=303)
