from src.model.doc import Doc
from src.model.block import Block


class Retriever:

    def __init__(self, db_client, plan_doc: Doc, content_doc: Doc, content_fr_doc: Doc, collection_name: str):
        plan_blocks: [Block] = plan_doc.blocks
        content_blocks: [Block] = content_doc.blocks
        content_fr_blocks: [Block] = content_fr_doc.blocks
        for pb, cb in zip(plan_blocks, content_blocks):
            cb.specials = pb.specials
        for cb, cb_fr in zip(content_blocks, content_fr_blocks):
            cb.content_fr = cb_fr.content
            cb.title_fr = cb_fr.title
        self.collection = db_client.create_collection(name=collection_name)
        self.collection.add(
            documents=[block.content for block in plan_blocks],
            ids=[block.index for block in plan_blocks],
            metadatas=[block.to_dict() for block in content_blocks]
        )

    def similarity_search(self, query: str) -> {}:
        res = self.collection.query(query_texts=query)
        block_dict_sources = res['metadatas'][0]
        distances = res['distances'][0]
        blocks = []
        for bd, d in zip(block_dict_sources, distances):
            b = Block().from_dict(bd)
            b.distance = d
            blocks.append(b)
        return blocks