app.py

"""
Chatty - Ethics Chat Application
A Flask web application that provides a chat interface for the findEthics-Atlas API.
"""

from flask import Flask, render_template, request, jsonify, redirect, url_for, flash, session, make_response
from flask_login import login_user, logout_user, login_required, current_user
from flask_wtf.csrf import CSRFProtect
import requests
import json
import os
import logging
from datetime import datetime, timedelta
from collections import defaultdict
import time
import jwt

from config import get_config, validate_environment
from auth import init_login_manager, authenticate_user, create_user_account
from models import User, ChatSession
from functools import wraps
import secrets
import string

# Validate configuration before starting
if not validate_environment():
    print("Configuration validation failed. Please check your environment variables.")
    exit(1)

app = Flask(__name__)

# Load configuration based on environment
config_class = get_config()
app.config.from_object(config_class)

# CSRF Protection with conditional enabling
if app.config.get('WTF_CSRF_ENABLED', True):
    csrf = CSRFProtect(app)
    app.logger.info("CSRF protection enabled")
else:
    csrf = None
    app.logger.info("CSRF protection disabled for this environment")

# Add CSRF error handler for debugging
@app.errorhandler(400)
def handle_csrf_error(e):
    """Handle CSRF errors with better debugging"""
    error_description = str(e.description) if hasattr(e, 'description') else str(e)
    
    # Log CSRF errors for debugging
    app.logger.warning(f"CSRF Error: {error_description}")
    app.logger.warning(f"Request headers: {dict(request.headers)}")
    app.logger.warning(f"Request form: {dict(request.form)}")
    app.logger.warning(f"Session: {dict(session)}")
    
    if 'CSRF' in error_description:
        flash('Security token expired. Please try again.', 'error')
        # Redirect back to the same page to regenerate CSRF token
        if request.endpoint in ['login', 'register']:
            return redirect(url_for(request.endpoint))
        return redirect(url_for('login'))
    
    return render_template('errors/400.html'), 400

# Configure logging
logging.basicConfig(
    level=getattr(logging, app.config['LOG_LEVEL']),
    format='%(asctime)s %(levelname)s %(name)s %(message)s',
    filename=app.config.get('LOG_FILE')
)

# Initialize Flask-Login
init_login_manager(app)

# Stateless authentication functions for Hugging Face Spaces
def create_auth_token(user_id, email):
    """Create a JWT token for stateless authentication"""
    payload = {
        'user_id': str(user_id),
        'email': email,
        'exp': datetime.utcnow() + timedelta(hours=48),  # 48 hour expiry
        'iat': datetime.utcnow()
    }
    return jwt.encode(payload, app.config['SECRET_KEY'], algorithm='HS256')

def verify_auth_token(token):
    """Verify and decode JWT token"""
    try:
        payload = jwt.decode(token, app.config['SECRET_KEY'], algorithms=['HS256'])
        return payload
    except jwt.ExpiredSignatureError:
        app.logger.info("Auth token expired")
        return None
    except jwt.InvalidTokenError:
        app.logger.info("Invalid auth token")
        return None

def get_current_user_from_token():
    """Get current user from JWT token in cookie"""
    token = request.cookies.get('auth_token')
    
    # Check if token is missing or explicitly deleted
    if not token or token == 'deleted':
        return None
    
    payload = verify_auth_token(token)
    if not payload:
        return None
    
    try:
        user = User.get_by_id(payload['user_id'])
        if user and user.email == payload['email']:
            return user
    except Exception as e:
        app.logger.warning(f"Failed to load user from token: {e}")
    
    return None

# Custom login_required decorator for Hugging Face Spaces load balancing
def hf_login_required(f):
    """
    Custom login_required decorator that uses stateless JWT authentication
    for Hugging Face Spaces load balancing compatibility
    """
    @wraps(f)
    def decorated_function(*args, **kwargs):
        # First check if user is already authenticated via Flask-Login
        if current_user.is_authenticated:
            return f(*args, **kwargs)
        
        # Try to authenticate using JWT token from cookie
        user = get_current_user_from_token()
        if user:
            # Restore the user session for this request
            login_user(user, remember=True)
            app.logger.info(f"Restored user from JWT token: {user.email}")
            return f(*args, **kwargs)
        
        # If all else fails, redirect to login
        app.logger.info(f"Authentication required - redirecting to login")
        app.logger.info(f"Session data: {dict(session)}")
        app.logger.info(f"Auth token present: {'auth_token' in request.cookies}")
        app.logger.info(f"Request IP: {request.environ.get('HTTP_X_FORWARDED_FOR', request.remote_addr)}")
        
        return redirect(url_for('login', next=request.url))
    
    return decorated_function

# Initialize database connection
from database import init_database
try:
    init_database()
except Exception as e:
    print(f"Warning: Database initialization failed: {e}")

# Rate limiting for login attempts
login_attempts = defaultdict(list)
MAX_LOGIN_ATTEMPTS = app.config['MAX_LOGIN_ATTEMPTS']
RATE_LIMIT_WINDOW = app.config['RATE_LIMIT_WINDOW']

def is_rate_limited(ip_address):
    """Check if IP address is rate limited for login attempts"""
    current_time = time.time()
    attempts = login_attempts[ip_address]
    
    # Remove attempts older than the rate limit window
    login_attempts[ip_address] = [attempt_time for attempt_time in attempts 
                                 if current_time - attempt_time < RATE_LIMIT_WINDOW]
    
    # Check if current attempts exceed the limit
    return len(login_attempts[ip_address]) >= MAX_LOGIN_ATTEMPTS

def record_login_attempt(ip_address):
    """Record a failed login attempt"""
    login_attempts[ip_address].append(time.time())

# Configuration
API_URL = app.config['API_URL']
TIMEOUT = app.config['API_TIMEOUT']

# Anonymous user session management with JWT tokens for stateless operation
def generate_anonymous_id():
    """Generate a unique anonymous user ID"""
    random_string = ''.join(secrets.choice(string.ascii_lowercase + string.digits) for _ in range(8))
    return f"anon_{random_string}"

def create_anonymous_token(anonymous_id, chat_history=None, message_count=0, first_message_time=None):
    """Create a JWT token for anonymous users"""
    payload = {
        'anonymous_id': anonymous_id,
        'is_anonymous': True,
        'exp': datetime.utcnow() + timedelta(hours=6),  # 6 hour expiry
        'iat': datetime.utcnow(),
        'chat_history': chat_history or [],
        'message_count': message_count,
        'first_message_time': first_message_time
    }
    return jwt.encode(payload, app.config['SECRET_KEY'], algorithm='HS256')

def verify_anonymous_token(token):
    """Verify and decode anonymous JWT token"""
    try:
        payload = jwt.decode(token, app.config['SECRET_KEY'], algorithms=['HS256'])
        if payload.get('is_anonymous'):
            return payload
        return None
    except jwt.ExpiredSignatureError:
        app.logger.info("Anonymous token expired")
        return None
    except jwt.InvalidTokenError:
        app.logger.info("Invalid anonymous token")
        return None

def get_or_create_anonymous_session():
    """Get existing anonymous session from token or create new one"""
    token = request.cookies.get('anon_token')
    
    # Check if token exists and is not deleted
    if token and token != 'deleted':
        payload = verify_anonymous_token(token)
        if payload:
            # Restore session from token
            session['is_anonymous'] = True
            session['anonymous_id'] = payload.get('anonymous_id')
            session['anonymous_chat_history'] = payload.get('chat_history', [])
            session['anonymous_message_count'] = payload.get('message_count', 0)
            session['anonymous_first_message_time'] = payload.get('first_message_time')
            app.logger.info(f"Restored anonymous session from token: {payload.get('anonymous_id')}")
            return payload.get('anonymous_id'), token
    
    # Create new anonymous session
    anonymous_id = generate_anonymous_id()
    session.permanent = True
    session['is_anonymous'] = True
    session['anonymous_id'] = anonymous_id
    session['anonymous_chat_history'] = []
    session['anonymous_message_count'] = 0
    session['anonymous_first_message_time'] = None
    
    # Create JWT token
    token = create_anonymous_token(anonymous_id)
    app.logger.info(f"Created new anonymous session: {anonymous_id}")
    return anonymous_id, token

def update_anonymous_token():
    """Update anonymous JWT token with current session data"""
    if not session.get('is_anonymous'):
        return None
    
    # Keep only last 20 messages in token
    chat_history = session.get('anonymous_chat_history', [])
    if len(chat_history) > 20:
        chat_history = chat_history[-20:]
    
    return create_anonymous_token(
        anonymous_id=session.get('anonymous_id'),
        chat_history=chat_history,
        message_count=session.get('anonymous_message_count', 0),
        first_message_time=session.get('anonymous_first_message_time')
    )

def get_anonymous_chat_history():
    """Get chat history from session for anonymous user"""
    return session.get('anonymous_chat_history', [])

def add_anonymous_message(message, response):
    """Add message and response to anonymous session history"""
    if 'anonymous_chat_history' not in session:
        session['anonymous_chat_history'] = []
    
    session['anonymous_chat_history'].append({
        'user': message,
        'assistant': response
    })
    
    # Keep only last 20 messages for anonymous users
    if len(session['anonymous_chat_history']) > 20:
        session['anonymous_chat_history'] = session['anonymous_chat_history'][-20:]

def check_anonymous_rate_limit():
    """Check if anonymous user has exceeded rate limit"""
    current_time = time.time()
    
    # Initialize rate limit tracking
    if 'anonymous_first_message_time' not in session or session.get('anonymous_first_message_time') is None:
        session['anonymous_first_message_time'] = current_time
        session['anonymous_message_count'] = 0
    
    # Check if rate limit window has passed
    time_since_first = current_time - session['anonymous_first_message_time']
    if time_since_first > app.config['ANONYMOUS_RATE_LIMIT_WINDOW']:
        # Reset rate limit window
        session['anonymous_first_message_time'] = current_time
        session['anonymous_message_count'] = 0
    
    # Check if limit exceeded
    if session.get('anonymous_message_count', 0) >= app.config['ANONYMOUS_RATE_LIMIT']:
        time_remaining = app.config['ANONYMOUS_RATE_LIMIT_WINDOW'] - time_since_first
        return False, f"Rate limit exceeded. Please wait {int(time_remaining/60)} minutes."
    
    # Increment message count
    session['anonymous_message_count'] = session.get('anonymous_message_count', 0) + 1
    return True, None

def make_api_request(message, history=None, user_id=None):
    """
    Make POST request to the findEthics-Atlas API
    """
    try:
        # Prepare the request payload
        payload = {
            "prompt": message,
            "history": history or []
        }
        
        # Add user_id if provided (for authenticated users)
        if user_id:
            payload["user_id"] = str(user_id)
        
        # Make the API request
        response = requests.post(
            API_URL,
            json=payload,
            timeout=TIMEOUT,
            headers={
                "Content-Type": "application/json"
            }
        )
        
        # Check if request was successful
        response.raise_for_status()
        
        # Parse the response
        data = response.json()
        return {
            "success": True,
            "response": data.get("response", ""),
            "error": None
        }
        
    except requests.exceptions.Timeout:
        return {
            "success": False,
            "response": "",
            "error": "Request timed out. Please try again."
        }
    except requests.exceptions.RequestException as e:
        return {
            "success": False,
            "response": "",
            "error": f"API request failed: {str(e)}"
        }
    except json.JSONDecodeError:
        return {
            "success": False,
            "response": "",
            "error": "Invalid response from API"
        }
    except Exception as e:
        return {
            "success": False,
            "response": "",
            "error": f"Unexpected error: {str(e)}"
        }

@app.route('/')
def index():
    """
    Main route - show anonymous chat by default, authenticated chat if logged in
    """
    if current_user.is_authenticated:
        # Authenticated users get the full chat interface
        return render_template('index.html', is_anonymous=False)
    
    # Check if anonymous access is enabled
    if not app.config.get('ANONYMOUS_ENABLED', True):
        flash('Please login or register to use the chat.', 'info')
        return redirect(url_for('login'))
    
    # Initialize or restore anonymous session with JWT token
    anonymous_id, token = get_or_create_anonymous_session()
    
    # Create response with anonymous chat interface
    response = make_response(render_template('index.html', is_anonymous=True))
    
    # Set anonymous JWT token cookie
    response.set_cookie(
        'anon_token',
        token,
        max_age=6*60*60,  # 6 hours
        secure=False,      # HF handles HTTPS at proxy level
        httponly=False,    # Allow JS access for updates
        samesite=None      # Most permissive for cross-origin
    )
    
    return response

@app.route('/register', methods=['GET', 'POST'])
def register():
    """
    Registration route - display form and handle user registration
    """
    # If user is already logged in, redirect to chat
    if current_user.is_authenticated:
        return redirect(url_for('chat_interface'))
    
    if request.method == 'GET':
        # Display registration form
        return render_template('register.html')
    
    elif request.method == 'POST':
        # Handle registration form submission
        try:
            # Get form data
            email = request.form.get('email', '').strip()
            password = request.form.get('password', '')
            confirm_password = request.form.get('confirm_password', '')
            
            # Validate and create user account
            success, user, errors = create_user_account(email, password, confirm_password)
            
            if success:
                # Registration successful
                flash('Registration successful! Please log in with your credentials.', 'success')
                return redirect(url_for('login'))
            else:
                # Registration failed - display errors
                for error in errors:
                    flash(error, 'error')
                return render_template('register.html', email=email)
                
        except Exception as e:
            flash('Registration failed. Please try again.', 'error')
            return render_template('register.html')

@app.route('/login', methods=['GET', 'POST'])
def login():
    """
    Login route - display form and handle user authentication
    """
    # If user is already logged in, redirect to chat
    if current_user.is_authenticated:
        return redirect(url_for('chat_interface'))
    
    if request.method == 'GET':
        # Display login form
        return render_template('login.html')
    
    elif request.method == 'POST':
        # Handle login form submission
        try:
            # Get client IP address
            client_ip = request.environ.get('HTTP_X_FORWARDED_FOR', request.remote_addr)
            if client_ip:
                client_ip = client_ip.split(',')[0].strip()
            
            # Check rate limiting
            if is_rate_limited(client_ip):
                flash('Too many login attempts. Please try again in 15 minutes.', 'error')
                return render_template('login.html'), 429
            
            # Get form data
            email = request.form.get('email', '').strip()
            password = request.form.get('password', '')
            
            # Authenticate user
            success, user, error_message = authenticate_user(email, password)
            
            if success and user:
                # Login successful - create session and JWT token
                session.permanent = True
                login_user(user, remember=True)
                
                # Create JWT token for stateless authentication across load-balanced servers
                auth_token = create_auth_token(user.id, user.email)
                
                # Store user info in session as backup
                session['user_id'] = str(user.id)
                session['user_email'] = user.email
                session['login_time'] = datetime.utcnow().isoformat()
                
                # Debug information
                app.logger.info(f"Session after login_user: {dict(session)}")
                app.logger.info(f"Current user authenticated: {current_user.is_authenticated}")
                app.logger.info(f"Current user ID: {getattr(current_user, 'id', 'None')}")
                app.logger.info(f"JWT token created: {bool(auth_token)}")
                app.logger.info(f"Request IP: {request.environ.get('HTTP_X_FORWARDED_FOR', request.remote_addr)}")
                
                flash('Login successful!', 'success')
                
                # Create response with JWT token cookie
                next_page = request.args.get('next')
                if next_page and next_page.startswith('/'):
                    response = make_response(redirect(next_page))
                else:
                    response = make_response(redirect(url_for('chat_interface')))
                
                # Set JWT token cookie with very permissive settings for HF Spaces
                response.set_cookie(
                    'auth_token',
                    auth_token,
                    max_age=48*60*60,  # 48 hours
                    secure=False,      # HF handles HTTPS at proxy level
                    httponly=False,    # Allow JS access for iframe compatibility
                    samesite=None      # Most permissive for cross-origin
                )
                
                return response
            else:
                # Login failed - record attempt and display error
                record_login_attempt(client_ip)
                flash(error_message or 'Invalid email or password', 'error')
                return render_template('login.html', email=email)
                
        except Exception as e:
            flash('Login failed. Please try again.', 'error')
            return render_template('login.html')

@app.route('/logout', methods=['POST'])
def logout():
    """
    Logout route - terminate user session and redirect to login
    """
    try:
        # Debug logging before logout
        app.logger.info(f"Logout attempt for user: {current_user.email if current_user.is_authenticated else 'unknown'}")
        
        # Clear session data first
        session.clear()
        
        # Log out the current user (Flask-Login)
        logout_user()
        
        # Set a flag to prevent auto-login on next request
        session['just_logged_out'] = True
        
        # Flash success message
        flash('You have been logged out successfully.', 'info')
        
        # Create response and clear all cookies properly
        response = make_response(redirect(url_for('index')))
        
        # Clear auth_token with matching parameters from login
        response.set_cookie(
            'auth_token', 
            value='deleted',  # Set to 'deleted' instead of empty
            max_age=0,  # Expire immediately
            expires=0,  # Also set expires for compatibility
            path='/',
            domain=None,
            secure=False,  # Match the setting from login
            httponly=False,  # Match the setting from login
            samesite=None  # Match the setting from login
        )
        
        # Clear anonymous token
        response.set_cookie(
            'anon_token',
            value='deleted',
            max_age=0,
            expires=0,
            path='/',
            domain=None,
            secure=False,
            httponly=False,
            samesite=None
        )
        
        # Clear Flask session cookie
        response.set_cookie(
            'session',
            value='deleted',
            max_age=0,
            expires=0,
            path='/',
            domain=None,
            secure=False,
            httponly=True,  # Session cookie should be httponly
            samesite='Lax'  # Flask default
        )
        
        # Debug logging after logout
        app.logger.info(f"Logout completed. All cookies cleared. Redirecting to index.")
        
        return response
        
    except Exception as e:
        app.logger.error(f"Error during logout: {e}")
        # Even if there's an error, still try to log out
        session.clear()
        logout_user()
        flash('Logout completed.', 'info')
        
        # Clear cookies even on error
        response = make_response(redirect(url_for('index')))
        response.set_cookie('auth_token', value='deleted', max_age=0, expires=0, path='/', 
                          secure=False, httponly=False, samesite=None)
        response.set_cookie('anon_token', value='deleted', max_age=0, expires=0, path='/', 
                          secure=False, httponly=False, samesite=None)
        response.set_cookie('session', value='deleted', max_age=0, expires=0, path='/', 
                          secure=False, httponly=True, samesite='Lax')
        
        return response


@app.route('/chat')
@hf_login_required
def chat_interface():
    """
    Protected chat interface - requires authentication
    Redirects to homepage which handles both anonymous and authenticated chat
    """
    return redirect(url_for('index'))

@app.route('/api/chat', methods=['POST'])
def chat():
    """
    API route to handle chat requests - supports both authenticated and anonymous users
    """
    try:
        # Check if user is anonymous - try to restore from token first
        is_anonymous = False
        anonymous_token = None
        
        if not current_user.is_authenticated:
            # Try to get anonymous session from token
            token = request.cookies.get('anon_token')
            app.logger.debug(f"Anonymous token present: {bool(token)}")
            
            if token and token != 'deleted':
                payload = verify_anonymous_token(token)
                app.logger.debug(f"Anonymous token verification result: {bool(payload)}")
                
                if payload:
                    # Restore session from token
                    session['is_anonymous'] = True
                    session['anonymous_id'] = payload.get('anonymous_id')
                    session['anonymous_chat_history'] = payload.get('chat_history', [])
                    session['anonymous_message_count'] = payload.get('message_count', 0)
                    session['anonymous_first_message_time'] = payload.get('first_message_time')
                    is_anonymous = True
                    anonymous_token = token
                    app.logger.debug(f"Anonymous session restored for ID: {payload.get('anonymous_id')}")
                else:
                    # Token verification failed, create new anonymous session
                    app.logger.debug("Anonymous token verification failed, creating new session")
                    anonymous_id, new_token = get_or_create_anonymous_session()
                    is_anonymous = True
                    anonymous_token = new_token
            else:
                # No token found, create new anonymous session
                app.logger.debug("No anonymous token found, creating new session")
                anonymous_id, new_token = get_or_create_anonymous_session()
                is_anonymous = True
                anonymous_token = new_token
            
            # If still not anonymous, check session (fallback)
            if not is_anonymous:
                is_anonymous = session.get('is_anonymous', False)
                app.logger.debug(f"Using session fallback, is_anonymous: {is_anonymous}")
        
        # If not anonymous and not authenticated, require login
        if not is_anonymous and not current_user.is_authenticated:
            app.logger.warning(f"Unauthorized access to /api/chat - is_anonymous: {is_anonymous}, authenticated: {current_user.is_authenticated}")
            return jsonify({
                "success": False,
                "error": "Your session has expired. Please log in again.",
                "redirect": url_for('login')
            }), 401
        
        data = request.get_json()
        
        if not data or 'message' not in data:
            return jsonify({
                "success": False,
                "error": "Please provide a message to send."
            }), 400
        
        message = data['message'].strip()
        if not message:
            return jsonify({
                "success": False,
                "error": "Your message cannot be empty. Please type something."
            }), 400
        
        # Check message length (different limits for anonymous vs authenticated)
        max_length = app.config['ANONYMOUS_MAX_MESSAGE_LENGTH'] if is_anonymous else 5000
        if len(message) > max_length:
            return jsonify({
                "success": False,
                "error": f"Your message is too long. Please keep it under {max_length} characters."
            }), 400
        
        # Check rate limit for anonymous users
        if is_anonymous:
            allowed, error_msg = check_anonymous_rate_limit()
            if not allowed:
                return jsonify({
                    "success": False,
                    "error": error_msg
                }), 429
        
        # Get chat history for context
        atlas_history = []
        
        try:
            if is_anonymous:
                # For anonymous users, get history from session and format for Atlas API
                session_history = get_anonymous_chat_history()
                
                # Convert anonymous session history to Atlas API format
                for exchange in session_history[-app.config['CONVERSATION_HISTORY_LIMIT']:]:
                    if exchange.get("user") and exchange.get("assistant"):
                        atlas_history.extend([
                            {"role": "user", "content": exchange.get("user", "")},
                            {"role": "assistant", "content": exchange.get("assistant", "")}
                        ])
                
                app.logger.debug(f"Anonymous chat with {len(atlas_history)} history messages")
            else:
                # For authenticated users, get conversation history from database
                atlas_history = ChatSession.get_conversation_history_for_atlas(
                    user_id=current_user.id,
                    limit=app.config['CONVERSATION_HISTORY_LIMIT'],
                    max_tokens=app.config['MAX_HISTORY_TOKENS']
                )
                
                app.logger.debug(f"Authenticated chat for user {current_user.id} with {len(atlas_history)} history messages")
                
        except Exception as e:
            # If history retrieval fails, log the error but continue with empty history
            app.logger.warning(f"Failed to retrieve conversation history: {e}")
            atlas_history = []
        
        # Make API request with history (falls back to empty history if retrieval failed)
        if is_anonymous:
            # For anonymous users, don't pass user_id
            result = make_api_request(message, atlas_history)
        else:
            # For authenticated users, pass the user_id to the API
            result = make_api_request(message, atlas_history, user_id=str(current_user.id))
        
        # If successful, save the message and response
        if result.get('success'):
            response_text = result.get('response', '')
            
            if is_anonymous:
                # Save to session for anonymous users
                add_anonymous_message(message, response_text)
                
                # Update the anonymous JWT token with new data
                updated_token = update_anonymous_token()
                
                result['history_saved'] = False  # Indicate temporary storage
                result['user_id'] = session.get('anonymous_id', 'anonymous')
                result['is_anonymous'] = True
                result['message_count'] = session.get('anonymous_message_count', 0)
                result['messages_remaining'] = max(0, app.config['ANONYMOUS_RATE_LIMIT'] - session.get('anonymous_message_count', 0))
                
                # Include updated token in response for client to store
                if updated_token:
                    result['update_token'] = updated_token
            else:
                # Save to database for authenticated users
                session_metadata = {
                    'timestamp': datetime.utcnow().isoformat(),
                    'api_url': API_URL,
                    'user_agent': request.headers.get('User-Agent', ''),
                    'ip_address': request.environ.get('HTTP_X_FORWARDED_FOR', request.remote_addr)
                }
                
                save_success = ChatSession.save_message(
                    user_id=current_user.id,
                    message=message,
                    response=response_text,
                    session_metadata=session_metadata
                )
                
                if not save_success:
                    # Log the error but don't fail the request
                    app.logger.warning(f"Failed to save chat message for user {current_user.id}")
                    # Optionally notify user that history wasn't saved
                    result['history_saved'] = False
                else:
                    result['history_saved'] = True
                
                result['user_id'] = current_user.id
                result['is_anonymous'] = False
        else:
            # Enhance error messages from API
            error_msg = result.get('error', 'Unknown error occurred')
            if 'timeout' in error_msg.lower():
                result['error'] = "The request took too long to process. Please try again."
            elif 'connection' in error_msg.lower():
                result['error'] = "Unable to connect to the chat service. Please check your internet connection and try again."
            elif 'api' in error_msg.lower():
                result['error'] = "The chat service is temporarily unavailable. Please try again in a few moments."
        
        # If we have an updated anonymous token, set it in the response cookie
        if is_anonymous:
            if 'update_token' in result:
                response = make_response(jsonify(result))
                response.set_cookie(
                    'anon_token',
                    result['update_token'],
                    max_age=6*60*60,  # 6 hours
                    secure=False,      # HF handles HTTPS
                    httponly=False,    # Allow JS access
                    samesite=None
                )
                # Remove token from JSON response to avoid duplication
                del result['update_token']
                response.data = json.dumps(result)
                response.content_type = 'application/json'
                return response
            elif anonymous_token:
                # Ensure the anonymous token is set in the response
                response = make_response(jsonify(result))
                response.set_cookie(
                    'anon_token',
                    anonymous_token,
                    max_age=6*60*60,  # 6 hours
                    secure=False,      # HF handles HTTPS
                    httponly=False,    # Allow JS access
                    samesite=None
                )
                return response
        
        return jsonify(result)
        
    except Exception as e:
        user_identifier = session.get('anonymous_id', 'unknown') if is_anonymous else (current_user.id if current_user.is_authenticated else 'unknown')
        app.logger.error(f"Chat API error for user {user_identifier}: {str(e)}")
        return jsonify({
            "success": False,
            "error": "We're experiencing technical difficulties. Please try again in a few moments."
        }), 500

@app.route('/api/user-history')
@hf_login_required
def get_user_history():
    """
    API route to retrieve authenticated user's chat history
    """
    try:
        # Check if user session is still valid
        if not current_user.is_authenticated:
            return jsonify({
                "success": False,
                "error": "Your session has expired. Please log in again.",
                "redirect": url_for('login')
            }), 401
        
        # Get limit from query parameters (default 50)
        limit = request.args.get('limit', 50, type=int)
        limit = max(1, min(limit, 100))  # Ensure between 1 and 100
        
        # Get user's chat history
        history = ChatSession.get_user_history(current_user.id, limit=limit)
        
        return jsonify({
            "success": True,
            "history": history,
            "count": len(history),
            "user_id": current_user.id
        })
        
    except Exception as e:
        app.logger.error(f"Error retrieving user history for user {current_user.id if current_user.is_authenticated else 'unknown'}: {str(e)}")
        return jsonify({
            "success": False,
            "error": "Unable to load your chat history. Please refresh the page and try again."
        }), 500

@app.route('/api/clear-chat', methods=['POST'])
def clear_chat():
    """
    API route to clear user's chat history
    """
    try:
        # Check if user is authenticated
        is_authenticated = current_user.is_authenticated if hasattr(current_user, 'is_authenticated') else False
        
        if is_authenticated:
            # Clear authenticated user's chat history
            result = ChatSession.clear_user_history(current_user.id)
            if result:
                app.logger.info(f"Cleared chat history for user {current_user.id}")
                return jsonify({
                    "success": True,
                    "message": "Chat history cleared successfully"
                })
            else:
                return jsonify({
                    "success": False,
                    "error": "Failed to clear chat history"
                }), 500
        else:
            # For anonymous users, clear session data
            session.pop('anonymous_chat_history', None)
            session['anonymous_message_count'] = 0  # Reset message counter
            app.logger.info(f"Cleared anonymous chat history for session {session.get('anonymous_id', 'unknown')}")
            return jsonify({
                "success": True,
                "message": "Chat history cleared successfully"
            })
    
    except Exception as e:
        user_identifier = current_user.id if current_user.is_authenticated else session.get('anonymous_id', 'unknown')
        app.logger.error(f"Error clearing chat history for user {user_identifier}: {str(e)}")
        return jsonify({
            "success": False,
            "error": "Unable to clear chat history. Please try again."
        }), 500

@app.route('/health')
def health():
    """
    Health check endpoint
    """
    return jsonify({"status": "healthy"})


# Error handler for 400 is already defined above with CSRF debugging

@app.errorhandler(401)
def handle_unauthorized(e):
    """Handle unauthorized access"""
    if request.is_json:
        return jsonify({
            'success': False,
            'error': 'Session expired. Please log in again.',
            'redirect': url_for('login')
        }), 401
    return render_template('errors/401.html'), 401

@app.errorhandler(429)
def handle_rate_limit(e):
    """Handle rate limiting errors"""
    return render_template('errors/429.html'), 429

@app.errorhandler(500)
def handle_server_error(e):
    """Handle internal server errors"""
    app.logger.error(f'Server Error: {e}')
    return render_template('errors/500.html'), 500

@app.before_request
def check_session_expiry():
    """Check for session expiry and handle gracefully"""
    # Skip this check for logout route to prevent re-authentication
    if request.endpoint == 'logout':
        return
    
    # Check if user just logged out
    if session.get('just_logged_out'):
        # Keep the flag for one more request to ensure cookies are cleared
        if request.endpoint == 'index':
            session.pop('just_logged_out', None)
        # Don't restore from token if just logged out
        return
    
    # List of endpoints that should not trigger auto-login
    no_auto_login_endpoints = ['login', 'register', 'static', 'health', 'index', 'api.chat']
    
    # Try to restore user from JWT token if not authenticated
    if not current_user.is_authenticated and request.endpoint not in no_auto_login_endpoints:
        # Check if auth_token cookie is 'deleted' - if so, don't try to restore
        auth_token = request.cookies.get('auth_token')
        if auth_token and auth_token != 'deleted':
            user = get_current_user_from_token()
            if user:
                login_user(user, remember=True)
                app.logger.info(f"Auto-restored user from JWT token: {user.email}")
    
    if current_user.is_authenticated:
        # Check if session is about to expire (within 1 hour)
        if hasattr(session, 'permanent') and session.permanent:
            from datetime import datetime, timedelta
            # Track session start time
            if 'session_start_time' not in session:
                session['session_start_time'] = datetime.utcnow().isoformat()
            
            session_start = session.get('session_start_time')
            if session_start:
                if isinstance(session_start, str):
                    try:
                        session_start = datetime.fromisoformat(session_start)
                    except:
                        session_start = datetime.utcnow()
                        session['session_start_time'] = session_start.isoformat()
                elif not isinstance(session_start, datetime):
                    session_start = datetime.utcnow()
                    session['session_start_time'] = session_start.isoformat()
                
                time_since_start = datetime.utcnow() - session_start
                if time_since_start > timedelta(hours=23):  # Warn 1 hour before expiry
                    flash('Your session will expire soon. Please save your work.', 'warning')

if __name__ == "__main__":
    # For Hugging Face Spaces deployment
    import os
    port = int(os.environ.get("PORT", 7860))
    app.run(debug=False, host='0.0.0.0', port=port)