import streamlit as st
from langchain.embeddings import HuggingFaceEmbeddings
from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate
import chromadb

# CONFIG
BASE_IP = "192.168.0.49"
LM_STUDIO_PORT = "1234"
CHROMA_PATH = "vector_db"
MODEL_NAME = "lmstudio-community/qwen/qwen3-next-80b-a3b-instruct-q8_0.gguf"  # Use "llama3", "phi3", etc.
EMBEDDING_MODEL = "all-MiniLM-L6-v2"

# Load embedding model
embedder = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)

# Load local LLM for answering
llm = Ollama(model=MODEL_NAME, temperature=0.3)

# Initialize Chroma client
client = chromadb.PersistentClient(path=CHROMA_PATH)
collection = client.get_collection("documents")

# Prompt template
prompt_template = """
You are a helpful assistant that answers questions using ONLY the context provided.
Do not make up information or use external knowledge.

Question: {question}

Context:
{context}

If you cannot find an answer, say "I don't know based on the provided documents."

Answer:
"""

prompt = PromptTemplate.from_template(prompt_template)

# Streamlit UI
st.title("📄 Local RAG Knowledge Assistant")
st.write("Upload files to `documents/` and run `ingest.py` first.")

query = st.text_input("Ask a question about your documents:",
                      placeholder="What are the key financial metrics?")

if query:
    with st.spinner("Searching for relevant info..."):
        # Embed query
        query_embedding = embedder.embed_query(query)

        # Retrieve top 5 most similar chunks
        results = collection.query(
            query_embeddings=[query_embedding],
            n_results=5,
            include=["documents", "metadatas"]
        )

        documents = results["documents"][0]
        metadatas = results["metadatas"][0]

        # Build context from retrieved chunks + metadata
        context = ""
        for i, doc in enumerate(documents):
            meta = metadatas[i]
            synopsis = meta.get("synopsis", "No summary")
            tags = ", ".join(meta.get("tags", [])) if isinstance(
                meta.get("tags"), list) else str(meta.get("tags"))
            source = meta.get("source", "Unknown")

            context += f"""
--- Document Snippet ---
{doc}

Synopsis: {synopsis}
Tags: {tags}
Source: {source}
---
"""

        # Ask LLM
        full_prompt = prompt.format(question=query, context=context)

    with st.spinner("Generating answer..."):
        response = llm.invoke(full_prompt)

    st.subheader("🔍 Answer:")
    st.write(response)

    st.subheader("📚 Sources (retrieved chunks):")
    for i, doc in enumerate(documents):
        meta = metadatas[i]
        source = meta.get("source", "Unknown")
        tags = ", ".join(meta.get("tags", [])) if isinstance(
            meta.get("tags"), list) else str(meta.get("tags"))
        st.markdown(f"**Source**: `{source}` | **Tags**: {tags}")
        st.text_area(f"Snippet {i+1}", doc, height=120, disabled=True)
