import streamlit as st
import os
from llama_index.core import VectorStoreIndex, Settings, SimpleDirectoryReader, StorageContext
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
# This is the new, correct import for the Gemini model from the modern library
from llama_index.llms.google import Gemini 

# --- App Configuration ---
st.set_page_config(
    page_title="Genshin Miliastra Sandbox Helper",
    page_icon="🤖",
    layout="centered",
    initial_sidebar_state="auto",
)

st.title("🤖 Genshin Miliastra Sandbox AI Helper")
st.caption("Ask me anything about how to use the Miliastra Wonderland sandbox!")

# --- Model and API Key Setup ---
with st.sidebar:
    st.header("Configuration")
    google_api_key = st.text_input("Enter your Google API Key", type="password")
    st.markdown(
        "Get your free API key from [Google AI Studio](https://aistudio.google.com/)."
    )

# --- Define paths for your raw data and the persistent storage ---
DATA_DIR = "./data"
STORAGE_DIR = "./storage"

# --- Caching Functions to improve performance ---

@st.cache_resource(show_spinner="Loading embedding model...")
def load_embed_model():
    """Loads and caches the embedding model from Hugging Face."""
    return HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

@st.cache_resource(show_spinner="Loading knowledge base... (This may take a moment on first startup)")
def load_index():
    """
    Loads or builds the vector index. 
    On the very first run, it will build the index from the documents in DATA_DIR.
    On subsequent runs, it will load the pre-built index from STORAGE_DIR.
    """
    Settings.embed_model = load_embed_model()

    if not os.path.exists(STORAGE_DIR):
        st.info("First time startup: Building the knowledge base from your documents. Please wait.")
        # Load all your .txt files from the 'data' directory
        documents = SimpleDirectoryReader(DATA_DIR).load_data()
        
        # Create the index from the documents
        index = VectorStoreIndex.from_documents(documents)
        
        # Save the newly built index to disk for future use
        index.storage_context.persist(persist_dir=STORAGE_DIR)
        st.success("Knowledge base built successfully!")
    else:
        # Load the existing index from the 'storage' directory
        storage_context = StorageContext.from_defaults(persist_dir=STORAGE_DIR)
        index = VectorStoreIndex.from_storage(storage_context)
        
    return index

# --- Main Application Logic ---
index = load_index()

if "messages" not in st.session_state:
    st.session_state.messages = [
        {"role": "assistant", "content": "Hi! How can I help you build in the Miliastra Sandbox today?"}
    ]

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

# Handle user input
if prompt := st.chat_input("Your question"):
    # Ensure the user has entered their API key
    if not google_api_key:
        st.info("Please add your Google API Key to continue.")
        st.stop()

    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)

    # Generate and display the assistant's response
    with st.chat_message("assistant"):
        message_placeholder = st.empty()
        message_placeholder.markdown("Thinking...")

        try:
            # Configure the LLM with the API key
            Settings.llm = Gemini(model_name="models/gemini-1.5-pro-latest", api_key=google_api_key)
            
            # Create a query engine from the loaded index
            query_engine = index.as_query_engine(similarity_top_k=3) # Get top 3 relevant text chunks
            response = query_engine.query(prompt)
            
            # Display the final answer
            message_placeholder.markdown(str(response))
            st.session_state.messages.append({"role": "assistant", "content": str(response)})

        except Exception as e:
            error_message = f"An error occurred: {e}. Please check your API key and try again."
            message_placeholder.error(error_message)
            st.session_state.messages.append({"role": "assistant", "content": error_message})