"""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', 'auth_redirect',
           '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, Response
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,
    require_subscription: bool = False,
    subscription_redirect: str = None,
    grace_period_days: int = 3,
):
    """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)
    
    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):
        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 = {}
            
            # 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 RedirectResponse(subscription_redirect, status_code=303)
                        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 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 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 46
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)
