import os
import json
import streamlit as st
from dotenv import load_dotenv
import google.generativeai as genai

# Load environment variables from .env file
load_dotenv()
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# Configure the Gemini API
genai.configure(api_key=GOOGLE_API_KEY)

# Directory to store session files
SESSION_DIR = "sessions"
os.makedirs(SESSION_DIR, exist_ok=True)

# Initialize session state
if "session_name" not in st.session_state:
    st.session_state.session_name = None
if "messages" not in st.session_state:
    st.session_state.messages = []

# Function to load session from file
def load_session(session_name):
    try:
        with open(os.path.join(SESSION_DIR, f"{session_name}.json"), "r") as f:
            return json.load(f)
    except FileNotFoundError:
        return []

# Function to save session to file
def save_session(session_name, messages):
    with open(os.path.join(SESSION_DIR, f"{session_name}.json"), "w") as f:
        json.dump(messages, f)

# Function to generate a session name using AI
def generate_session_name(prompt):
    model = genai.GenerativeModel("gemini-2.0-flash")
    response = model.generate_content(f"Suggest a unique and descriptive session name based on this message: {prompt}")
    return response.text.strip()

# Sidebar for session management
st.sidebar.title("Chat Sessions")

# List existing sessions
existing_sessions = [f[:-5] for f in os.listdir(SESSION_DIR) if f.endswith(".json")]

# Select session
selected_session = st.sidebar.selectbox("Select a session", [""] + existing_sessions)

if selected_session:
    st.session_state.session_name = selected_session
    st.session_state.messages = load_session(selected_session)
else:
    st.session_state.messages = []

# Create new session
new_session_name = st.sidebar.text_input("New session name")
if st.sidebar.button("Create Session"):
    if new_session_name:
        st.session_state.session_name = new_session_name
        st.session_state.messages = []
        save_session(new_session_name, [])
        st.rerun()

# Delete current session
if st.sidebar.button("Delete Session"):
    if st.session_state.session_name:
        os.remove(os.path.join(SESSION_DIR, f"{st.session_state.session_name}.json"))
        st.session_state.session_name = None
        st.session_state.messages = []
        st.rerun()

# Main chat interface
st.title("Gemini Chatbot with Session Management")

# Display chat history
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# Chat input
if prompt := st.chat_input("Enter your message:"):
    # Display user message
    st.chat_message("user").markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})

    # Generate and display AI response
    model = genai.GenerativeModel("gemini-2.0-flash")
    try:
        response = model.generate_content(prompt)
        st.chat_message("assistant").markdown(response.text)
        st.session_state.messages.append({"role": "assistant", "content": response.text})

        # Auto-generate session name if not set
        if not st.session_state.session_name:
            session_name = generate_session_name(prompt)
            st.session_state.session_name = session_name
            st.sidebar.write(f"Auto-generated session name: {session_name}")

        # Save session
        save_session(st.session_state.session_name, st.session_state.messages)
    except Exception as e:
        st.error(f"An error occurred: {e}")