import tiktoken
from random import shuffle

# `tokenizer`: used to count how many tokens
tokenizer_name = tiktoken.encoding_for_model('gpt-4')
tokenizer = tiktoken.get_encoding(tokenizer_name.name)

def tiktoken_len(text):
    # evaluate how many tokens for the given text
    tokens = tokenizer.encode(text, disallowed_special=())
    return len(tokens)


class Knowledge:
    def __init__(self, db):
        self.db = db
        self.contents = []

    def collect_knowledge(self, keywords_dict, max_query):
        """
        keywords_dict:
            {"machine learning": 5, "language model": 2};
        """
        db = self.db
        if max_query > 0:
            for kw in keywords_dict:
                docs = db.similarity_search_with_score(kw, k=max_query)
                for i in range(max_query):
                    content = {"content": docs[i][0].page_content.replace('\n', ' '),
                               "score": docs[i][1]}  # todo: add more meta information; clean the page_content
                    self.contents.append(content)
            # sort contents by score / shuffle
            shuffle(self.contents)

    def to_prompts(self, max_tokens=2048):
        if len(self.contents) == 0:
            return ""
        prompts = []
        tokens = 0
        for idx, content in enumerate(self.contents):
            prompt = "Reference {}: {}\n".format(idx, content["content"])
            tokens += tiktoken_len(prompt)
            if tokens >= max_tokens:
                break
            else:
                prompts.append(prompt)
        return "".join(prompts)