components/session_manager.py

import streamlit as st
from streamlit_js_eval import get_cookie, set_cookie
import uuid
from datetime import datetime
import re
from db.helpers import get_user_history, ensure_user_exists

def initialize_session_state():
    """Initialize session state with persistence for authenticated users"""
    
    # Check if we have authentication cookies first
    auth_cookie = get_cookie("tailored_auth")
    
    # Authentication defaults
    auth_defaults = {
        "authenticated": False,
        "username": "",
        "fullname": "",
        "current_page": "auth",
    }
    
    # If we have an auth cookie, restore authentication
    if auth_cookie and not st.session_state.get("authenticated"):
        try:
            import json
            auth_data = json.loads(auth_cookie)
            st.session_state.authenticated = True
            st.session_state.user_id = auth_data.get("user_id")
            st.session_state.username = auth_data.get("username")
            st.session_state.fullname = auth_data.get("fullname")
            st.session_state.current_page = "generator"
            print(f"✅ Restored auth from cookie for user: {st.session_state.username}")
        except Exception as e:
            print(f"❌ Error restoring auth from cookie: {e}")
            # Clear invalid cookie
            set_cookie("tailored_auth", "", expires=0)
    
    # Set authentication defaults if not already set
    for key, value in auth_defaults.items():
        if key not in st.session_state:
            st.session_state[key] = value
    
    # Only set app defaults if user is authenticated
    if st.session_state.authenticated:
        # Application defaults
        app_defaults = {
            "user_type": "",
            "original_prompt": "",
            "generated_output": "",
            "feedback_given": False,
            "regenerated": False,
            "content_source": "",
            "student_level": "",
            "pdf_export_data": None,
            "tutor_topic": "",                    
            "tutor_content_type": "",             
            "feedback_clarity": 3,
            "feedback_depth": 3,
            "feedback_complexity": "Just right",
            "feedback_comments": "",
            "original_filename": "content.pdf",
            "saved_to_history": False,
            "user_history": [],  
            "from_history": False,
            "showing_regeneration_prompt": False,
            "pending_model_switch": None,
            "previous_model": None,
            "regenerate_with_new_model": False,
            "scrolled_to_top": False,
            "regeneration_count": 0,  
            "regeneration_type": None,  
            "previous_feedback_given": False,
            "pending_regeneration": False,  
            "show_adaptation_message": True,
            "selected_model": "groq",  # Add default model selection
        }
        
        for key, value in app_defaults.items():
            if key not in st.session_state:
                st.session_state[key] = value
        
        # Load user history
        load_user_history_from_db()

def set_auth_cookie(user_data):
    """Set authentication cookie for persistence"""
    try:
        import json
        import time
        cookie_data = json.dumps({
            "user_id": user_data["user_id"],
            "username": user_data["username"],
            "fullname": user_data["fullname"],
            "timestamp": time.time()
        })
        # Set cookie to expire in 7 days
        set_cookie("tailored_auth", cookie_data, expires=7)
        print("✅ Auth cookie set")
    except Exception as e:
        print(f"❌ Error setting auth cookie: {e}")

def clear_auth_cookie():
    """Clear authentication cookie on logout"""
    try:
        set_cookie("tailored_auth", "", expires=0)
        print("✅ Auth cookie cleared")
    except Exception as e:
        print(f"❌ Error clearing auth cookie: {e}")

def load_user_history_from_db():
    """Load user's history from database into session state - ENHANCED ERROR HANDLING"""
    try:
        # Check if user_id exists before loading history
        if not hasattr(st.session_state, 'user_id') or not st.session_state.user_id:
            print("⚠️ No user_id found - skipping history load")
            st.session_state.user_history = []
            return
            
        user_id = st.session_state.user_id
        history = get_user_history(user_id)
        st.session_state.user_history = history
        print(f"✅ Loaded {len(history)} history entries for user {user_id[:8]}...")
        
    except Exception as e:
        print(f"❌ Error loading history from database: {e}")
        st.session_state.user_history = []

# def clear_session():
#     """Clear session state but preserve user identity and history"""
#     preserved_keys = ['user_id', 'current_page']  # REMOVED 'user_history' from preserved keys
#     preserved = {k: st.session_state[k] for k in preserved_keys if k in st.session_state}
    
#     st.session_state.clear()
    
#     # Restore preserved keys
#     for k, v in preserved.items():
#         st.session_state[k] = v
    
#     # Re-initialize defaults AND reload history
#     initialize_session_state()

def clear_session():
    """Clear session state but preserve user identity, auth, and navigation"""
    preserved_keys = [
        'user_id', 'username', 'fullname', 'authenticated', 
        'current_page', 'selected_model', 'user_history'
    ]
    
    preserved = {k: st.session_state[k] for k in preserved_keys if k in st.session_state}
    
    # Clear only content-related session state, not auth
    keys_to_clear = [
        'user_type', 'original_prompt', 'generated_output', 'feedback_given',
        'regenerated', 'content_source', 'student_level', 'pdf_export_data',
        'tutor_topic', 'tutor_content_type', 'feedback_clarity', 'feedback_depth',
        'feedback_complexity', 'feedback_comments', 'original_filename',
        'saved_to_history', 'from_history', 'showing_regeneration_prompt',
        'pending_model_switch', 'previous_model', 'regenerate_with_new_model',
        'scrolled_to_top', 'regeneration_count', 'regeneration_type',
        'previous_feedback_given', 'pending_regeneration', 'show_adaptation_message'
    ]
    
    for key in keys_to_clear:
        if key in st.session_state:
            del st.session_state[key]
    
    # Restore preserved keys
    for k, v in preserved.items():
        st.session_state[k] = v
    
    # Ensure current page is generator for new content
    st.session_state.current_page = "generator"
    
    print("✅ Session cleared for new content, auth preserved")

def update_session_state(**kwargs):
    # Update all the session state values
    for key, value in kwargs.items():
        st.session_state[key] = value
    
    # Auto-save to database if we have new content
    if (st.session_state.get("generated_output") and 
        not st.session_state.get("saved_to_history", False) and 
        not st.session_state.get("regenerated", False)):
        
        save_current_to_history()

def save_current_to_history():
    """Save current content to database and update session state"""
    from db.helpers import save_content_to_history
    
    try:
        # ✅ ADD THIS SAFETY CHECK
        if not st.session_state.get("user_id"):
            print("⚠️ No user_id available - skipping history save")
            return None
            
        if (st.session_state.generated_output and 
            not st.session_state.saved_to_history and 
            not st.session_state.regenerated):
            
            # Handle PDF data - ensure it's base64 string
            pdf_data = st.session_state.pdf_export_data
            if isinstance(pdf_data, bytes):
                import base64
                pdf_base64 = base64.b64encode(pdf_data).decode('utf-8')
            elif pdf_data is None:
                pdf_base64 = ""  # Handle case where PDF might be None
            else:
                pdf_base64 = pdf_data  # Assume it's already base64 string
            
            content_data = {
                "user_id": st.session_state.user_id,
                "user_type": st.session_state.user_type,
                "student_level": st.session_state.student_level,
                "topic": st.session_state.tutor_topic if st.session_state.user_type == "tutor" else "Simplified Content",
                "content_type": st.session_state.tutor_content_type if st.session_state.user_type == "tutor" else "Simplified Explanation",
                "prompt": st.session_state.original_prompt,
                "output": st.session_state.generated_output,
                "pdf_base64": pdf_base64,
                "filename": generate_history_filename(),
                "feedback_given": st.session_state.feedback_given,
                "generated_model": st.session_state.get("generated_model", "groq")
            }
            
            entry_id = save_content_to_history(content_data)
            if entry_id:
                st.session_state.current_history_id = entry_id
                st.session_state.saved_to_history = True
                # Reload history from database to include new entry
                load_user_history_from_db()
                print(f"✅ Auto-saved content to history: {entry_id}")
                return entry_id
            else:
                print("❌ Failed to auto-save content to history")
                return None
                
        return None
        
    except Exception as e:
        print(f"❌ Error in save_current_to_history: {e}")
        return None
        

def generate_history_filename():
    timestamp = datetime.now().strftime("%Y%m%d_%H%M")
    if st.session_state.user_type == "student":
        level_clean = re.sub(r'[^a-zA-Z0-9]', '_', st.session_state.student_level)
        return f"student_content_{level_clean}_{timestamp}.pdf"
    else:
        topic_clean = re.sub(r'[^a-zA-Z0-9]', '_', st.session_state.tutor_topic)[:20]
        content_type_clean = st.session_state.tutor_content_type.replace(' ', '_')
        return f"{content_type_clean}_{topic_clean}_{timestamp}.pdf"

def get_session_info():
    return {
        "user_id": st.session_state.user_id,
        "user_type": st.session_state.user_type,
        "has_output": bool(st.session_state.generated_output),
        "feedback_given": st.session_state.feedback_given,
        "regenerated": st.session_state.regenerated,
        "current_page": st.session_state.current_page,
        "current_history_id": st.session_state.current_history_id,
        "history_entries": len(st.session_state.user_history)
    }

def prepare_for_model_regeneration():
    """Prepare session state for model regeneration - UPDATED tracking"""
    # Preserve the essential content generation context
    preserved_data = {
        'user_type': st.session_state.user_type,
        'student_level': st.session_state.student_level,
        'content_source': st.session_state.content_source,
        'original_prompt': st.session_state.original_prompt,
        # For tutor flow
        'tutor_topic': st.session_state.get('tutor_topic', ''),
        'tutor_content_type': st.session_state.get('tutor_content_type', ''),
        # For student flow  
        'original_filename': st.session_state.get('original_filename', 'content.pdf'),
    }
    
    # Clear generation outputs but keep context
    keys_to_clear = [
        'generated_output', 'pdf_export_data', 'feedback_given', 
        'regenerated', 'current_history_id', 'saved_to_history',
        'feedback_clarity', 'feedback_depth', 'feedback_complexity', 
        'feedback_comments', 'scrolled_to_top'
    ]
    
    for key in keys_to_clear:
        if key in st.session_state:
            del st.session_state[key]
    
    # Track regeneration
    st.session_state.regeneration_count = st.session_state.get('regeneration_count', 0) + 1
    st.session_state.regeneration_type = 'model_switch'
    
    # Restore preserved context
    for key, value in preserved_data.items():
        if value:  # Only restore if we have actual values
            st.session_state[key] = value
    
    st.session_state.regenerate_with_new_model = True