auth_manager.py

"""
Authentication Manager for ID Agents
=====================================

Handles user authentication, session management, and access control
for the ID Agents application running on HF Spaces.
"""

import gradio as gr
import time
import secrets
from typing import Dict, Optional, Tuple, Any
from auth_config import authenticate_user, authenticate_with_code, get_user_capabilities

class AuthManager:
    def __init__(self):
        self.active_sessions: Dict[str, Dict] = {}
        self.session_timeout = 3600  # 1 hour in seconds
        
    def create_session(self, user_info: Dict) -> str:
        """Create a new user session"""
        session_id = secrets.token_urlsafe(32)
        self.active_sessions[session_id] = {
            "user_info": user_info,
            "created_at": time.time(),
            "last_activity": time.time(),
            "capabilities": get_user_capabilities(user_info["access_level"])
        }
        return session_id
    
    def validate_session(self, session_id: str) -> Tuple[bool, Optional[Dict]]:
        """Validate a session and return user info if valid"""
        if not session_id or session_id not in self.active_sessions:
            return False, None
        
        session = self.active_sessions[session_id]
        current_time = time.time()
        
        # Check if session has expired
        if current_time - session["last_activity"] > self.session_timeout:
            del self.active_sessions[session_id]
            return False, None
        
        # Update last activity
        session["last_activity"] = current_time
        return True, session
    
    def logout_session(self, session_id: str) -> bool:
        """Logout a user session"""
        if session_id in self.active_sessions:
            del self.active_sessions[session_id]
            return True
        return False
    
    def get_session_info(self, session_id: str) -> Optional[Dict]:
        """Get session information"""
        is_valid, session = self.validate_session(session_id)
        return session if is_valid else None

# Global auth manager instance
auth_manager = AuthManager()

def create_auth_interface():
    """Create the authentication interface"""
    
    with gr.Blocks(title="ID Agents - Beta Access") as auth_interface:
        
        # Add custom CSS for the auth interface
        auth_interface.css = """
        :root {
          --auth-bg: #1a1f2e;
          --auth-panel: #23263a;
          --auth-accent: #00ffe7;
          --auth-accent2: #7f5cff;
          --auth-text: #ffffff;
          --auth-muted: #a0a5b8;
          --auth-success: #00ff88;
          --auth-error: #ff4757;
        }

        body, .gradio-container, .gr-block, .gr-app {
          background: var(--auth-bg) !important;
          color: var(--auth-text) !important;
        }

        .auth-container {
          max-width: 500px;
          margin: 2rem auto;
          padding: 2rem;
          background: linear-gradient(135deg, var(--auth-panel) 60%, #2a2f45 100%);
          border: 2px solid var(--auth-accent2);
          border-radius: 20px;
          box-shadow: 0 8px 40px rgba(127, 92, 255, 0.3);
        }

        .auth-title {
          text-align: center;
          font-size: 2rem;
          font-weight: 700;
          margin-bottom: 0.5rem;
          background: linear-gradient(90deg, var(--auth-accent), var(--auth-accent2));
          -webkit-background-clip: text;
          -webkit-text-fill-color: transparent;
          background-clip: text;
        }

        .auth-subtitle {
          text-align: center;
          color: var(--auth-muted);
          margin-bottom: 2rem;
          font-size: 1.1rem;
        }

        .auth-tabs {
          background: transparent !important;
          border: none !important;
        }

        .auth-form {
          background: rgba(255, 255, 255, 0.05);
          border-radius: 12px;
          padding: 1.5rem;
          margin: 1rem 0;
        }

        .auth-input {
          background: rgba(255, 255, 255, 0.1) !important;
          border: 1px solid var(--auth-accent2) !important;
          border-radius: 8px !important;
          color: var(--auth-text) !important;
          margin-bottom: 1rem;
        }

        .auth-button {
          background: linear-gradient(90deg, var(--auth-accent2), var(--auth-accent)) !important;
          color: var(--auth-bg) !important;
          border: none !important;
          border-radius: 8px !important;
          font-weight: 600 !important;
          padding: 12px 24px !important;
          width: 100% !important;
          margin-top: 1rem;
        }

        .auth-success {
          color: var(--auth-success) !important;
          background: rgba(0, 255, 136, 0.1);
          border: 1px solid var(--auth-success);
          border-radius: 8px;
          padding: 1rem;
          margin: 1rem 0;
        }

        .auth-error {
          color: var(--auth-error) !important;
          background: rgba(255, 71, 87, 0.1);
          border: 1px solid var(--auth-error);
          border-radius: 8px;
          padding: 1rem;
          margin: 1rem 0;
        }

        .auth-info {
          background: rgba(0, 255, 231, 0.1);
          border: 1px solid var(--auth-accent);
          border-radius: 8px;
          padding: 1rem;
          margin: 1rem 0;
          color: var(--auth-accent);
        }

        .user-info {
          background: rgba(127, 92, 255, 0.1);
          border: 1px solid var(--auth-accent2);
          border-radius: 8px;
          padding: 1rem;
          margin: 1rem 0;
        }
        """
        
        # Session state
        session_id = gr.State("")
        
        with gr.Column(elem_classes="auth-container"):
            gr.HTML('<div class="auth-title">🦠 ID Agents</div>')
            gr.HTML('<div class="auth-subtitle">Beta Testing Access Portal</div>')
            
            # Authentication status
            auth_status = gr.Markdown("", visible=False)
            user_info_display = gr.Markdown("", visible=False)
            
            # Login tabs
            with gr.Tabs(elem_classes="auth-tabs") as auth_tabs:
                
                # Username/Password Login
                with gr.Tab("👤 User Login", elem_classes="auth-form"):
                    gr.Markdown("### Sign in with your beta testing credentials")
                    username_input = gr.Textbox(
                        label="Username",
                        placeholder="Enter your username...",
                        elem_classes="auth-input"
                    )
                    password_input = gr.Textbox(
                        label="Password", 
                        type="password",
                        placeholder="Enter your password...",
                        elem_classes="auth-input"
                    )
                    login_button = gr.Button("🔑 Sign In", elem_classes="auth-button")
                    
                    gr.HTML("""
                    <div class="auth-info">
                        <strong>Beta Testing Accounts:</strong><br>
                        • dr_smith (password: idweek2025)<br>
                        • id_fellow (password: hello)<br>
                        • pharmacist (password: stewardship)<br>
                        • ipc_nurse (password: infection)<br>
                        • researcher (password: research)<br>
                        • admin (password: idagents2025)
                    </div>
                    """)
                
                # Invitation Code Login
                with gr.Tab("🎫 Invitation Code", elem_classes="auth-form"):
                    gr.Markdown("### Enter your invitation code")
                    invitation_input = gr.Textbox(
                        label="Invitation Code",
                        placeholder="Enter your invitation code...",
                        elem_classes="auth-input"
                    )
                    code_login_button = gr.Button("🎟️ Access with Code", elem_classes="auth-button")
                    
                    gr.HTML("""
                    <div class="auth-info">
                        <strong>Sample Invitation Codes:</strong><br>
                        • IDWEEK2025 (VIP Access)<br>
                        • BETA-FELLOW (Fellowship Program)<br>
                        • PHARM-STEW (Pharmacy Access)<br>
                        • IPC-NURSE (Infection Prevention)<br>
                        • RESEARCH-AI (Clinical Research)
                    </div>
                    """)
            
            # Main app placeholder (hidden until authenticated)
            app_container = gr.HTML("", visible=False)
            
            # Logout button (hidden until authenticated)
            logout_button = gr.Button("🚪 Logout", visible=False, elem_classes="auth-button")
        
        # Authentication functions
        def handle_login(username: str, password: str, current_session: str):
            """Handle username/password login"""
            if not username or not password:
                return (
                    gr.update(value="⚠️ Please enter both username and password.", visible=True, elem_classes="auth-error"),
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    current_session
                )
            
            success, user_info = authenticate_user(username, password)
            
            if success:
                session_id = auth_manager.create_session(user_info)
                user_display = f"""
                <div class="user-info">
                    <strong>✅ Welcome, {user_info['full_name']}!</strong><br>
                    Role: {user_info['role']}<br>
                    Access Level: {user_info['access_level']}<br>
                    Email: {user_info['email']}
                </div>
                """
                
                return (
                    gr.update(value="🎉 Authentication successful! Loading ID Agents...", visible=True, elem_classes="auth-success"),
                    gr.update(value=user_display, visible=True),
                    gr.update(visible=False),  # Hide auth tabs
                    gr.update(visible=True),   # Show app container
                    gr.update(visible=True),   # Show logout button
                    session_id
                )
            else:
                return (
                    gr.update(value="❌ Invalid username or password. Please try again.", visible=True, elem_classes="auth-error"),
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    current_session
                )
        
        def handle_invitation_login(invitation_code: str, current_session: str):
            """Handle invitation code login"""
            if not invitation_code:
                return (
                    gr.update(value="⚠️ Please enter an invitation code.", visible=True, elem_classes="auth-error"),
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    current_session
                )
            
            success, user_info = authenticate_with_code(invitation_code.upper())
            
            if success:
                session_id = auth_manager.create_session(user_info)
                user_display = f"""
                <div class="user-info">
                    <strong>✅ Welcome, {user_info['full_name']}!</strong><br>
                    Role: {user_info['role']}<br>
                    Access Level: {user_info['access_level']}<br>
                    Email: {user_info['email']}<br>
                    <em>Authenticated via invitation code</em>
                </div>
                """
                
                return (
                    gr.update(value="🎉 Authentication successful! Loading ID Agents...", visible=True, elem_classes="auth-success"),
                    gr.update(value=user_display, visible=True),
                    gr.update(visible=False),  # Hide auth tabs
                    gr.update(visible=True),   # Show app container
                    gr.update(visible=True),   # Show logout button
                    session_id
                )
            else:
                return (
                    gr.update(value="❌ Invalid or expired invitation code. Please check your code and try again.", visible=True, elem_classes="auth-error"),
                    gr.update(visible=False),
                    gr.update(visible=True),
                    gr.update(visible=False),
                    gr.update(visible=False),
                    current_session
                )
        
        def handle_logout(current_session: str):
            """Handle user logout"""
            if current_session:
                auth_manager.logout_session(current_session)
            
            return (
                gr.update(value="", visible=False),
                gr.update(visible=False),
                gr.update(visible=True),   # Show auth tabs
                gr.update(visible=False),  # Hide app container
                gr.update(visible=False),  # Hide logout button
                "",  # Clear session
                "",  # Clear username
                "",  # Clear password
                ""   # Clear invitation code
            )
        
        # Wire up the authentication handlers
        login_button.click(
            fn=handle_login,
            inputs=[username_input, password_input, session_id],
            outputs=[auth_status, user_info_display, auth_tabs, app_container, logout_button, session_id]
        )
        
        code_login_button.click(
            fn=handle_invitation_login,
            inputs=[invitation_input, session_id],
            outputs=[auth_status, user_info_display, auth_tabs, app_container, logout_button, session_id]
        )
        
        logout_button.click(
            fn=handle_logout,
            inputs=[session_id],
            outputs=[auth_status, user_info_display, auth_tabs, app_container, logout_button, session_id, username_input, password_input, invitation_input]
        )
    
    return auth_interface, session_id

def check_user_permission(session_id: str, required_capability: str) -> Tuple[bool, Optional[str]]:
    """
    Check if user has permission for a specific capability
    
    Returns:
        (has_permission: bool, error_message: str or None)
    """
    session_info = auth_manager.get_session_info(session_id)
    
    if not session_info:
        return False, "Please log in to access this feature."
    
    capabilities = session_info["capabilities"]
    
    if required_capability not in capabilities or not capabilities[required_capability]:
        user_level = session_info["user_info"]["access_level"]
        return False, f"Your access level ({user_level}) doesn't permit this action. Contact admin for upgrade."
    
    return True, None

def get_user_limits(session_id: str) -> Dict[str, Any]:
    """Get user-specific limits based on their access level"""
    session_info = auth_manager.get_session_info(session_id)
    
    if not session_info:
        return {
            "max_agents": 0,
            "max_file_size_mb": 0,
            "can_upload_files": False
        }
    
    return {
        "max_agents": session_info["capabilities"]["max_agents"],
        "max_file_size_mb": session_info["capabilities"]["max_file_size_mb"],
        "can_upload_files": session_info["capabilities"]["can_upload_files"],
        "can_download_configs": session_info["capabilities"]["can_download_configs"],
        "can_see_debug_info": session_info["capabilities"]["can_see_debug_info"]
    }