from rag_model import get_qa_chain, HuggingFaceEmbeddings, Chroma
import os

def main():
    PERSIST_DIRECTORY = "chroma_db"
    
    if not os.path.exists(PERSIST_DIRECTORY):
        print("Knowledge base not built. Building now...")
        from rag_model import build_rag_system
        vectorstore = build_rag_system()
    else:
        print("Loading existing knowledge base...")
        embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
        vectorstore = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embeddings)

    qa_chain = get_qa_chain(vectorstore)
    
    if not qa_chain:
        print("Failed to initialize QA chain. Check your GOOGLE_API_KEY in .env")
        return

    print("\n--- Retail Product Knowledge Assistant ---")
    print("Type 'exit' to quit.")
    
    while True:
        query = input("\nAsk a question about our products: ")
        if query.lower() == 'exit':
            break
            
        print("Thinking...")
        try:
            result = qa_chain({"query": query})
            print(f"\nAnswer: {result['result']}")
            print("\nSources (Top Reviews):")
            seen_sources = set()
            for doc in result["source_documents"]:
                source_info = f"{doc.metadata['name']} (Brand: {doc.metadata['brand']})"
                if source_info not in seen_sources:
                    print(f"- {source_info}")
                    seen_sources.add(source_info)
        except Exception as e:
            print(f"An error occurred: {e}")

if __name__ == "__main__":
    main()