import os
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct

class VectorDB:
    def __init__(self, collection_name="wifi_lab", vector_size=4096):
        host = os.getenv("QDRANT_HOST", "localhost")
        port = int(os.getenv("QDRANT_PORT", 6333))
        self.client = QdrantClient(host=host, port=port)
        self.collection_name = collection_name
        
        # Создаем коллекцию, если её нет (размерность 4096 для Qwen 8B Embedding)
        self.client.recreate_collection(
            collection_name=self.collection_name,
            vectors_config=VectorParams(size=vector_size, distance=Distance.COSINE),
        )

    def add_packet(self, pkt_id, vector, metadata, text):
        self.client.upsert(
            collection_name=self.collection_name,
            points=[
                PointStruct(
                    id=pkt_id,
                    vector=vector,
                    payload={"metadata": metadata, "text": text}
                )
            ]
        )

    def find_similar(self, vector):
        return self.client.search(
            collection_name=self.collection_name,
            query_vector=vector,
            limit=3,
            with_payload=True
        )