from qdrant_client.http import models
import pickle as pickle
import torch
import io

device_str = "cuda:0" if torch.cuda.is_available() else "cpu"
device = torch.device(device_str)


class Device_Unpickler(pickle.Unpickler):

    def find_class(self, module, name):
        if module == "torch.storage" and name == "_load_from_bytes":
            return lambda b: torch.load(io.BytesIO(b), map_location=device_str)
        else:
            return super().find_class(module, name)


def pickle_to_document_store(path):
    with open(path, "rb") as f:
        document_store = Device_Unpickler(f).load()
    document_store.embeddings.encode_kwargs["device"] = device_str
    return document_store


def get_qdrant_filters(filter_dict: dict):
    """Build a Qdrant filter based on a filter dict.

    Filter dict must use metadata fields and be formated like:

    filter_dict = {'file_name':['file1', 'file2'],'sub_type':['text']}
    """
    return models.Filter(
        must=[
            models.FieldCondition(
                key=f"metadata.{field}",
                match=models.MatchAny(any=filter_dict[field]),
            )
            for field in filter_dict
        ]
    )