MyFastHtml/src/myfasthtml/auth/routes.py

"""
Authentication routes for FastHTML application.

Provides routes for login, register, logout, and protected pages.
"""

from fasthtml.common import RedirectResponse
from myauth import create_auth_app_for_sqlite

from ..auth.pages.LoginPage import LoginPage
from ..auth.pages.RegisterPage import RegisterPage
from ..auth.pages.WelcomePage import WelcomePage
from ..auth.utils import (
  login_user,
  register_user,
  logout_user,
  get_user_info
)


def setup_auth_routes(app, rt, mount_auth_app=True, sqlite_db_path="Users.db", base_url=None):
  """
  Setup all authentication and protected routes.

  Args:
      app: FastHTML application instance
      rt: Route decorator from FastHTML
      mount_auth_app: Whether to mount the auth FastApi API routes
      sqlite_db_path: by default, create a new SQLite database at this path
      base_url: Base URL for the application (default to localhost:5001 if not provided)
  """
  
  # ============================================================================
  # PUBLIC ROUTES (Login & Register)
  # ============================================================================
  
  @rt("/login", methods=["GET"])
  def get(error: str = None):
    """
    Display login page.

    Args:
        error: Optional error message from query params

    Returns:
        LoginPage component
    """
    return LoginPage(error_message=error)
  
  @rt("/login", methods=["POST"])
  def post(email: str, password: str, session, redirect_url: str = "/"):
    """
    Handle login form submission.

    Args:
        email: User email from form
        password: User password from form
        session: FastHTML session object
        redirect_url: URL to redirect after successful login

    Returns:
        RedirectResponse on success, or LoginPage with error on failure
    """
    # Attempt login
    auth_data = login_user(email, password, base_url=base_url)
    
    if auth_data:
      # Login successful - store tokens in session
      session['access_token'] = auth_data['access_token']
      session['refresh_token'] = auth_data['refresh_token']
      
      # Get user info and store in session
      user_info = get_user_info(auth_data['access_token'], base_url=base_url)
      if user_info:
        session['user_info'] = user_info
      
      # Redirect to protected page
      return RedirectResponse(redirect_url, status_code=303)
    else:
      # Login failed - return error message via HTMX
      return LoginPage(error_message="Invalid email or password. Please try again.")
  
  @rt("/register", methods=["GET"])
  def get(error: str = None):
    """
    Display registration page.

    Args:
        error: Optional error message from query params

    Returns:
        RegisterPage component
    """
    return RegisterPage(error_message=error)
  
  @rt("/register", methods=["POST"])
  def post(email: str, username: str, password: str, confirm_password: str, session):
    """
    Handle registration form submission.

    Args:
        email: User email from form
        username: User name of the
        password: User password from form
        confirm_password: Password confirmation from form
        session: FastHTML session object

    Returns:
        RegisterPage with success/error message via HTMX
    """
    # Validate password confirmation
    if password != confirm_password:
      return RegisterPage(error_message="Passwords do not match. Please try again.")
    
    # Validate password length
    if len(password) < 8:
      return RegisterPage(error_message="Password must be at least 8 characters long.")
    
    # Attempt registration
    result = register_user(email, username, password, base_url=base_url)
    
    if result:
      # Registration successful - show success message and auto-login
      auth_data = login_user(email, password)
      
      if auth_data:
        # Store tokens in session
        session['access_token'] = auth_data['access_token']
        session['refresh_token'] = auth_data['refresh_token']
        
        # Get user info and store in session
        user_info = get_user_info(auth_data['access_token'])
        if user_info:
          session['user_info'] = user_info
        
        # Redirect to welcome page
        return RedirectResponse("/", status_code=303)
      else:
        # Auto-login failed, redirect to login page
        return RedirectResponse("/login", status_code=303)
    else:
      # Registration failed
      return RegisterPage(error_message="Registration failed. Email may already be in use.")
  
  # ============================================================================
  # PROTECTED ROUTES (Require authentication)
  # ============================================================================
  
  @rt("/welcome")
  def get(session, auth):
    """
    Alternative welcome page route (protected).

    Args:
        session: FastHTML session object
        auth: User auth info from request scope

    Returns:
        WelcomePage component
    """
    user_info = session.get('user_info', {})
    return WelcomePage(user_info=user_info)
  
  @rt("/logout")
  def post(session):
    """
    Handle logout request.

    Revokes refresh token and clears session.

    Args:
        session: FastHTML session object

    Returns:
        RedirectResponse to login page
    """
    # Get refresh token from session
    refresh_token = session.get('refresh_token')
    
    # Revoke refresh token on backend
    if refresh_token:
      logout_user(refresh_token)
    
    # Clear session
    session.clear()
    
    # Redirect to login page
    return RedirectResponse("/login", status_code=303)
  
  def mount_auth_fastapi_api():
    # Mount FastAPI auth backend
    auth_api = create_auth_app_for_sqlite(sqlite_db_path, "jwt-secret-to-change")
    app.mount("/auth", auth_api)
  
  if mount_auth_app:
    mount_auth_fastapi_api()