import streamlit as st
import os
import numpy as np
import faiss
from groq import Groq
from sentence_transformers import SentenceTransformer

# =========================================================
# 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  
✅ Groq API  
✅ Retrieval-Augmented Generation (RAG)

Runs on Hugging Face Spaces.
""")

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

    # =====================================================
    # LOAD EMBEDDING MODEL
    # =====================================================
    embedding_model = SentenceTransformer(
        "all-MiniLM-L6-v2"
    )

    # =====================================================
    # GROQ API KEY
    # =====================================================
    groq_api_key = os.getenv(
        "GROQ_API_KEY"
    )

    # If missing
    if not groq_api_key:

        st.error(
            "❌ GROQ_API_KEY is missing!"
        )

        st.stop()

    # Create Groq client
    client = Groq(
        api_key=groq_api_key
    )

    # =====================================================
    # DATASET PATH
    # =====================================================
    BASE_DIR = os.path.dirname(
        os.path.abspath(__file__)
    )

    data_path = os.path.join(
        BASE_DIR,
        "datasets"
    )

    # =====================================================
    # CHECK DATASET FOLDER
    # =====================================================
    if not os.path.exists(data_path):

        st.error(
            "❌ datasets folder not found!"
        )

        st.write(
            "Current files:",
            os.listdir(BASE_DIR)
        )

        st.stop()

    # =====================================================
    # READ TXT FILES
    # =====================================================
    all_texts = []

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

    # No TXT files
    if len(txt_files) == 0:

        st.error(
            "❌ No TXT files found!"
        )

        st.stop()

    # =====================================================
    # LOAD DATASET CONTENT
    # =====================================================
    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 TEXT
    # =====================================================
    full_text = " ".join(all_texts)

    # =====================================================
    # CHUNKING
    # =====================================================
    chunks = [
        full_text[i:i + 500]
        for i in range(
            0,
            len(full_text),
            500
        )
    ]

    # Optional optimization
    chunks = chunks[:2000]

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

    embeddings = np.array(
        embeddings,
        dtype=np.float32
    )

    # Normalize embeddings
    faiss.normalize_L2(
        embeddings
    )

    # =====================================================
    # CREATE FAISS INDEX
    # =====================================================
    dimension = embeddings.shape[1]

    index = faiss.IndexFlatIP(
        dimension
    )

    index.add(embeddings)

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

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

    (
        embedding_model,
        chunks,
        index,
        total_files,
        client
    ) = 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 query
    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:

            # =================================================
            # QUERY EMBEDDING
            # =================================================
            query_embedding = embedding_model.encode(
                [query]
            )

            query_embedding = np.array(
                query_embedding,
                dtype=np.float32
            )

            # Normalize query embedding
            faiss.normalize_L2(
                query_embedding
            )

            # =================================================
            # SEARCH FAISS
            # =================================================
            distances, indices = index.search(
                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 short
- Stay accurate
- If answer is not present, say:
"I don't know based on the provided context."
"""

            # =================================================
            # GROQ RESPONSE
            # =================================================
            response = client.chat.completions.create(
                model="llama-3.1-8b-instant",
                messages=[
                    {
                        "role": "system",
                        "content": (
                            "You are a Harry Potter expert assistant."
                        )
                    },
                    {
                        "role": "user",
                        "content": prompt
                    }
                ],
                temperature=0.2,
                max_tokens=200
            )

            answer = (
                response
                .choices[0]
                .message
                .content
            )

        except Exception as e:

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

    # =====================================================
    # SAVE ASSISTANT RESPONSE
    # =====================================================
    st.session_state.messages.append(
        {
            "role": "assistant",
            "content": answer
        }
    )

    # =====================================================
    # DISPLAY 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 + Groq + FAISS + SentenceTransformers
</center>
""",
    unsafe_allow_html=True
)