import os
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from langchain_community.document_loaders import TextLoader, PyPDFLoader, DirectoryLoader
from langchain.text_splitter import Language

class LocalKnowledgeBase:
    def __init__(self, collection_name="documents", persist_dir="./qdrant_data"):
        # Инициализируем клиент Qdrant
        self.client = QdrantClient(
            path=persist_dir,  # Локальное хранение
            prefer_grpc=True
        )
        
        self.collection_name = collection_name
        self.embeddings = embeddings
        
        # Создаем коллекцию если её нет
        self._create_collection()
    
    def _create_collection(self):
        try:
            self.client.get_collection(self.collection_name)
            print(f"Коллекция {self.collection_name} уже существует")
        except:
            # Создаем новую коллекцию
            self.client.create_collection(
                collection_name=self.collection_name,
                vectors_config=VectorParams(
                    size=768,  # Размерность эмбеддингов nomic-embed-text
                    distance=Distance.COSINE
                )
            )
            print(f"Создана коллекция {self.collection_name}")
    
    def load_documents(self, directory_path: str):
        """Загрузка документов из директории"""
        loaders = {
            '.pdf': PyPDFLoader,
            '.txt': lambda path: DirectoryLoader(path, glob="**/*.txt"),
            '.docx': lambda path: DirectoryLoader(path, glob="**/*.docx"),
            '.hpp': lambda path: TextLoader(path, encoding='utf-8'),
            '.cpp': lambda path: TextLoader(path, encoding='utf-8'),   # добавить
            '.h': lambda path: TextLoader(path, encoding='utf-8'),     # добавить
            '.cc': lambda path: TextLoader(path, encoding='utf-8'),    # добавить
            
        }
        
        all_documents = []
        
        for ext, loader_class in loaders.items():
            for file_path in os.listdir(directory_path):
                if file_path.endswith(ext):
                    full_path = os.path.join(directory_path, file_path)
                    try:
                        loader = loader_class(full_path)
                        documents = loader.load()
                        all_documents.extend(documents)
                        print(f"Загружен {file_path}: {len(documents)} страниц")
                    except Exception as e:
                        print(f"Ошибка загрузки {file_path}: {e}")
        
        # Разбиваем на чанки
        text_splitter = RecursiveCharacterTextSplitter.from_language(
            language=Language.CPP,
            chunk_size=1000,
            chunk_overlap=200,
            length_function=len
        )
        
        chunks = text_splitter.split_documents(all_documents)
        print(f"Всего чанков: {len(chunks)}")
        
        # Создаем эмбеддинги и сохраняем в Qdrant
        self._index_documents(chunks)
        
        return chunks
    
    def _index_documents(self, documents):
        """Индексация документов в Qdrant"""
        points = []
        
        for i, doc in enumerate(documents):
            # Создаем эмбеддинг для каждого чанка
            embedding = self.embeddings.embed_query(doc.page_content)
            
            point = PointStruct(
                id=i,
                vector=embedding,
                payload={
                    "text": doc.page_content,
                    "source": doc.metadata.get("source", "unknown"),
                    "page": doc.metadata.get("page", 0),
                    "file_type": "cpp"
                }
            )
            points.append(point)
            
            # Пакетная загрузка каждые 100 точек
            if len(points) >= 100:
                self.client.upsert(
                    collection_name=self.collection_name,
                    points=points
                )
                points = []
                print(f"Индексировано {i+1} документов")
        
        # Загружаем оставшиеся
        if points:
            self.client.upsert(
                collection_name=self.collection_name,
                points=points
            )
        
        print(f"Индексация завершена. Всего документов: {len(documents)}")
    
    def search(self, query: str, top_k: int = 5):
        """Поиск в базе знаний"""
        # Создаем эмбеддинг запроса
        query_embedding = self.embeddings.embed_query(query)
        
        # Ищем в Qdrant
        search_result = self.client.search(
            collection_name=self.collection_name,
            query_vector=query_embedding,
            limit=top_k
        )
        
        # Форматируем результаты
        results = []
        for hit in search_result:
            results.append({
                "text": hit.payload["text"],
                "score": hit.score,
                "source": hit.payload.get("source", "unknown")
            })
        
        return results

# Инициализируем базу знаний
kb = LocalKnowledgeBase()

# Загружаем документы (если есть)
if os.path.exists("./documents"):
    kb.load_documents("./documents")