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

# ==== Setup ====
st.set_page_config(page_title="🪖 US Army RAG Assistant", layout="wide")

# ==== Connect to Groq ====
groq_api_key = os.environ.get("api_key5")
client = Groq(api_key=groq_api_key)

# ==== Load FAISS index and metadata ====
index = faiss.read_index("army_qa_faiss.index")
with open("army_qa_metadata.pkl", "rb") as f:
    metadata = pickle.load(f)

encoder = SentenceTransformer("all-MiniLM-L6-v2")

# ==== RAG Function ====
def rag_query(user_query, top_k=5):
    query_embedding = encoder.encode([user_query])
    query_embedding = np.array(query_embedding).astype("float32")

    distances, indices = index.search(query_embedding, top_k)

    context_chunks = []
    for i in indices[0]:
        if 0 <= i < len(metadata):
            row = metadata[i]
            question = row["question"]
            context_chunks.append(f"Q: {question}\nA: [answer from database]")

    context = "\n---\n".join(context_chunks)

    prompt = f"""You are a highly trained assistant with access to US Army Field Manual knowledge.
Use the context below to answer the user's question.

Context:
{context}

User Question:
{user_query}

Answer:"""

    response = client.chat.completions.create(
        model="deepseek-r1-distill-llama-70b",
        messages=[{"role": "user", "content": prompt}]
    )

    return response.choices[0].message.content.strip()

# ==== Streamlit UI ====
st.title("🪖 US Army RAG Assistant")
st.markdown("Ask tactical questions based on **US Army Field Manual** knowledge.")
st.divider()

recommended_questions = [
    "What are the fundamentals of offensive operations?",
    "How does the Army define mission command?",
    "What is the purpose of a support zone?",
    "How do units prepare for degraded and denied environments (D3SOE)?",
    "Explain the role of reconnaissance in the offense."
]

st.markdown("💡 **Recommended Questions:**")
for q in recommended_questions:
    st.markdown(f"- {q}")

# Chat state
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# ==== Input ====
user_input = st.text_input("Your question", placeholder="Type your question here...")

if st.button("Submit") and user_input:
    if user_input.strip().lower() == "terminate":
        st.session_state.chat_history.append((user_input, "👋 Goodbye, soldier. Stay sharp!"))
    else:
        answer = rag_query(user_input)
        st.session_state.chat_history.append((user_input, answer))

if st.button("Clear Chat"):
    st.session_state.chat_history = []

# ==== Display Chat ====
for i, (user, bot) in enumerate(st.session_state.chat_history):
    st.markdown(f"**You:** {user}")
    st.markdown(f"**Assistant:** {bot}")
    st.divider()