"""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', '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', 'route_user_after_login',
           'require_tenant_access', 'handle_login_request', 'handle_oauth_callback', 'handle_logout']

# %% ../nbs/04_utils_auth.ipynb 2
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
import os
import uuid
import json
import logging
import time
from dotenv import load_dotenv

load_dotenv()

# Module-level logger - configured by app via configure_logging()
logger = logging.getLogger(__name__)

# %% ../nbs/04_utils_auth.ipynb 6
# 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 9
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 12
# 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 17
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,
):
    """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)
    
    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)
        
    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)
    """
    skip_patterns = []
    if include_defaults:
        skip_patterns.extend(DEFAULT_SKIP_AUTH)
    if skip:
        skip_patterns.extend(skip)
    
    def check_auth(req, sess):
        if session_key not in sess:
            return RedirectResponse(redirect_path, status_code=303)
        
        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 = {}
    
    return Beforeware(check_auth, skip=skip_patterns)

# %% ../nbs/04_utils_auth.ipynb 21
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 24
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 28
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 33
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 36
def create_user_session(session: dict, global_user: GlobalUser, membership: Membership):
    """Create authenticated session after successful OAuth login."""
    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()


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 41
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 45
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)
