import os
from langchain_chroma import Chroma
from langchain_community.embeddings import HuggingFaceEmbeddings
from huggingface_hub import InferenceClient

# --- Configuration ---
CHROMA_PATH = "./chroma_db"

# 1. Initialize DB and Embeddings ONCE globally
print("Loading embeddings and database...")
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
db = Chroma(persist_directory=CHROMA_PATH, embedding_function=embeddings)

# 2. Connect directly via Hugging Face's Native Client
hf_token = os.environ.get("HF_TOKEN")
if not hf_token:
    raise ValueError("HF_TOKEN missing. Please add your Hugging Face Access Token to the Space Secrets.")

# We use Zephyr through the modern InferenceClient to bypass the routing bug
client = InferenceClient(
    model="HuggingFaceH4/zephyr-7b-beta",
    token=hf_token
)

def get_rag_response(query_text):
    """This function is called by your Gradio app.py"""
    
    # Retrieve the relevant chunks from your database
    results = db.similarity_search(query_text, k=3)

    if len(results) == 0:
        return "I couldn't find any highly relevant notes to answer that."

    # Format the chunks into a single text block
    context_text = "\n\n---\n\n".join([doc.page_content for doc in results])

    # Build the prompt directly as a string
    prompt = f"""You are a helpful, private AI assistant. 
Answer the question based ONLY on the following context from my personal notes:

{context_text}

---

Question: {query_text}"""

    # Ask the model using the modern 'conversational' format!
    messages = [{"role": "user", "content": prompt}]
    
    # This completely replaces the broken LangChain llm.invoke()
    response = client.chat_completion(messages, max_tokens=512)
    response_text = response.choices[0].message.content

    # Format the final output to return to the Gradio UI
    final_answer = response_text.strip() + "\n\n### 📚 Sources Used:\n"
    for doc in results:
        source_name = doc.metadata.get('source', 'Unknown')
        final_answer += f"- **{source_name}**\n"

    return final_answer