import os
import argparse
from dotenv import load_dotenv
from langchain_openai import OpenAIEmbeddings
from langchain_qdrant import QdrantVectorStore
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_core.documents import Document
from qdrant_client import QdrantClient

# Загрузка переменных окружения
load_dotenv()

# Инициализация эмбеддингов
embeddings = OpenAIEmbeddings(
    model="text-embedding-3-small",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.getenv("OPENAI_API_KEY"),
)

# Инициализация Qdrant
client = QdrantClient(url="http://localhost:6333")
collection_name = "knowledge_base"
vector_store = QdrantVectorStore(
    client=client,
    collection_name=collection_name,
    embeddings=embeddings,
)

# Чанкинг
splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)

def load_directory(dir_path: str):
    """Загружает все .txt файлы из директории в векторную базу."""
    for root, dirs, files in os.walk(dir_path):
        for file in files:
            if file.lower().endswith(".txt"):
                path = os.path.join(root, file)
                with open(path, "r", encoding="utf-8") as f:
                    text = f.read()
                chunks = splitter.split_text(text)
                docs = [Document(page_content=chunk, metadata={"title": file}) for chunk in chunks]
                vector_store.add_documents(docs)
                print(f"Added {len(docs)} chunks from {file}")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Load documents into Qdrant.")
    parser.add_argument("directory", help="Path to directory with .txt files")
    args = parser.parse_args()
    load_directory(args.directory)