import streamlit as st
import os
import numpy as np
import faiss
from sentence_transformers import SentenceTransformer
from transformers import pipeline

# =========================================================
# PAGE CONFIG
# =========================================================
st.set_page_config(
    page_title="Harry Potter RAG Chatbot",
    page_icon="⚡",
    layout="wide"
)

# =========================================================
# CUSTOM CSS
# =========================================================
st.markdown("""
<style>

.main {
    background-color: #0E1117;
    color: white;
}

.chat-user {
    background-color: #1E293B;
    padding: 15px;
    border-radius: 12px;
    margin-bottom: 10px;
}

.chat-bot {
    background-color: #111827;
    padding: 15px;
    border-radius: 12px;
    margin-bottom: 10px;
}

</style>
""", unsafe_allow_html=True)

# =========================================================
# TITLE
# =========================================================
st.title("⚡ Harry Potter RAG Chatbot")
st.markdown("### Ask anything from the Harry Potter universe")

# =========================================================
# SIDEBAR
# =========================================================
with st.sidebar:

    st.header("⚙️ Settings")

    top_k = st.slider(
        "Retrieved Context Chunks",
        min_value=1,
        max_value=10,
        value=3
    )

    st.markdown("---")

    st.markdown("## 📚 About")

    st.write("""
This chatbot uses:

✅ Sentence Transformers  
✅ FAISS Vector Search  
✅ Hugging Face Transformers  
✅ Retrieval-Augmented Generation (RAG)

Runs on Hugging Face Spaces.
""")

# =========================================================
# LOAD LLM
# =========================================================
@st.cache_resource
def load_llm():

    pipe = pipeline(
        "text-generation",
        model="google/gemma-2b"
    )

    return pipe

pipe = load_llm()

# =========================================================
# LOAD EVERYTHING ONLY ONCE
# =========================================================
@st.cache_resource
def load_rag_system():

    # Load embedding model
    embedding_model = SentenceTransformer(
        "all-MiniLM-L6-v2"
    )

    # Dataset path
    data_path = "./src/dataset"

    # Check dataset folder
    if not os.path.exists(data_path):

        st.error("❌ dataset folder not found!")
        st.write("Current files:", os.listdir("./src"))
        st.stop()

    # Read txt files
    all_texts = []

    txt_files = [
        file for file in os.listdir(data_path)
        if file.endswith(".txt")
    ]

    if len(txt_files) == 0:

        st.error("❌ No TXT files found in dataset folder!")
        st.stop()

    # Load file contents
    for file_name in txt_files:

        file_path = os.path.join(data_path, file_name)

        with open(
            file_path,
            "r",
            encoding="utf-8"
        ) as f:

            content = f.read().strip()

            if content:
                all_texts.append(content)

    # Combine all text
    full_text = " ".join(all_texts)

    # Chunking
    chunks = [
        chunk.strip()
        for chunk in full_text.split(". ")
        if chunk.strip()
    ]

    # Speed optimization
    chunks = chunks[:2000]

    # Create embeddings
    embeddings = embedding_model.encode(
        chunks,
        show_progress_bar=True
    )

    # Create FAISS index
    dimension = embeddings.shape[1]

    index = faiss.IndexFlatL2(dimension)

    index.add(np.array(embeddings))

    return (
        embedding_model,
        chunks,
        index,
        len(txt_files)
    )

# =========================================================
# LOAD DATA
# =========================================================
with st.spinner("⚡ Loading AI model and dataset..."):

    embedding_model, chunks, index, total_files = load_rag_system()

st.success(
    f"✅ Loaded {total_files} dataset files successfully!"
)

# =========================================================
# SESSION STATE
# =========================================================
if "messages" not in st.session_state:

    st.session_state.messages = []

# =========================================================
# DISPLAY CHAT HISTORY
# =========================================================
for message in st.session_state.messages:

    if message["role"] == "user":

        st.markdown(
            f"""
<div class="chat-user">
<b>🧑 You:</b><br><br>
{message["content"]}
</div>
""",
            unsafe_allow_html=True
        )

    else:

        st.markdown(
            f"""
<div class="chat-bot">
<b>⚡ AI:</b><br><br>
{message["content"]}
</div>
""",
            unsafe_allow_html=True
        )

# =========================================================
# CHAT INPUT
# =========================================================
query = st.chat_input(
    "Ask a Harry Potter question..."
)

# =========================================================
# PROCESS QUERY
# =========================================================
if query:

    # Save user message
    st.session_state.messages.append(
        {
            "role": "user",
            "content": query
        }
    )

    # Display user message
    st.markdown(
        f"""
<div class="chat-user">
<b>🧑 You:</b><br><br>
{query}
</div>
""",
        unsafe_allow_html=True
    )

    # AI Processing
    with st.spinner("🔍 Searching Hogwarts Library..."):

        try:

            # Encode query
            query_embedding = embedding_model.encode([query])

            # Search FAISS
            distances, indices = index.search(
                np.array(query_embedding),
                k=top_k
            )

            # Retrieve chunks
            retrieved_chunks = [
                chunks[i]
                for i in indices[0]
            ]

            retrieved_text = "\n".join(retrieved_chunks)

            # Prompt
            prompt = f"""
You are a Harry Potter expert assistant.

Use ONLY the provided context.

================ CONTEXT ================
{retrieved_text}

================ QUESTION ================
{query}

Instructions:
- Give a clear answer
- Keep it beginner-friendly
- Keep it short and accurate
"""

            # Generate response
            response = pipe(
                prompt,
                max_new_tokens=200,
                do_sample=True
            )

            answer = response[0]["generated_text"]

        except Exception as e:

            answer = f"❌ Error: {str(e)}"

    # Save assistant response
    st.session_state.messages.append(
        {
            "role": "assistant",
            "content": answer
        }
    )

    # Display assistant response
    st.markdown(
        f"""
<div class="chat-bot">
<b>⚡ AI:</b><br><br>
{answer}
</div>
""",
        unsafe_allow_html=True
    )

    # Show retrieved context
    with st.expander("📚 Retrieved Context"):

        st.write(retrieved_text)

# =========================================================
# FOOTER
# =========================================================
st.markdown("---")

st.markdown(
    """
<center>
⚡ Built with Streamlit + Transformers + FAISS + SentenceTransformers
</center>
""",
    unsafe_allow_html=True
)