"""Thin Qdrant wrapper: one collection per embedding model under test."""

import os
import re
from uuid import uuid4

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


def collection_name_for(model_name):
    return "books__" + re.sub(r"[^a-zA-Z0-9]+", "_", model_name).strip("_")


def _dir_size_bytes(path):
    total = 0
    for root, _dirs, files in os.walk(path):
        for name in files:
            fp = os.path.join(root, name)
            if os.path.isfile(fp):
                total += os.path.getsize(fp)
    return total


class BookVectorStore:
    def __init__(self, qdrant_path):
        self.qdrant_path = qdrant_path
        self.client = QdrantClient(path=qdrant_path)

    def index(self, model_name, embeddings, books):
        name = collection_name_for(model_name)
        if self.client.collection_exists(name):
            self.client.delete_collection(name)

        self.client.create_collection(
            collection_name=name,
            vectors_config=VectorParams(
                size=embeddings.shape[1],
                distance=Distance.COSINE,
            ),
        )

        points = [
            PointStruct(
                id=str(uuid4()),
                vector=embedding.tolist(),
                payload={"book_id": book["book_id"]},
            )
            for book, embedding in zip(books, embeddings)
        ]
        self.client.upsert(collection_name=name, points=points)
        return name

    def search(self, model_name, query_embedding, top_k):
        name = collection_name_for(model_name)
        results = self.client.query_points(
            collection_name=name,
            query=query_embedding.tolist(),
            limit=top_k,
            with_payload=True,
        ).points
        return [point.payload["book_id"] for point in results]

    def collection_disk_size_mb(self, model_name):
        name = collection_name_for(model_name)
        collection_dir = os.path.join(self.qdrant_path, "collection", name)
        if not os.path.isdir(collection_dir):
            return None
        return _dir_size_bytes(collection_dir) / (1024 ** 2)