import os
from dotenv import load_dotenv
from agent import create_agent
from vectorstore import create_vectorstore, load_documents

# Load environment variables
load_dotenv()


def initialize_knowledge_base(documents_dir: str = "./documents"):
    """Initialize the vectorstore and load documents if not already loaded."""
    vectorstore = create_vectorstore()
    
    # Check if vectorstore is empty
    collection_count = vectorstore._collection.count()
    if collection_count == 0 and os.path.exists(documents_dir):
        print(f"Loading documents from {documents_dir}...")
        load_documents(documents_dir, vectorstore)
    else:
        print(f"Vectorstore already contains {collection_count} documents")
    
    return vectorstore


def main():
    """Main CLI chat loop."""
    print("=" * 50)
    print("RAG Agent with ChromaDB and Web Search")
    print("=" * 50)
    
    # Initialize knowledge base
    initialize_knowledge_base()
    
    # Create agent
    agent = create_agent()
    
    print("\nAgent ready. Type 'exit' to quit.\n")
    
    while True:
        try:
            query = input("Запрос: ").strip()
            
            if query.lower() == "exit":
                print("Goodbye!")
                break
            
            if not query:
                continue
            
            # Run agent
            result = agent.invoke({"input": query})
            answer = result.get("output", "No answer generated")
            
            print(f"\n{answer}\n")
            
        except KeyboardInterrupt:
            print("\nGoodbye!")
            break
        except Exception as e:
            print(f"Error: {e}")


if __name__ == "__main__":
    main()