import itertools
from typing import List

import tiktoken
from langchain.docstore.document import Document
from langchain.schema import AIMessage, BaseRetriever, HumanMessage

from scheduler.block import BaseBlock


class SlidingDocChain(BaseBlock):
    def __init__(
        self,
        retriever: BaseRetriever,
        max_tokens: int = 1024,
        link_key: str = "link",
        title_key: str = "title",
        idx_key: str = "idx",
    ):
        super(SlidingDocChain, self).__init__()
        self.enc = tiktoken.get_encoding("cl100k_base")
        self.retriever = retriever
        self.link_key = link_key
        self.max_tokens = max_tokens
        self.title_key = title_key
        self.idx_key = idx_key

    def __call__(self, text: str) -> str:
        docs = self.retriever.get_relevant_documents(text)
        for idx in range(len(docs)):
            if len(self.enc.encode(self.document(docs[: idx + 1]))) > self.max_tokens:
                return self.document(docs[:idx])
        return self.document(docs)

    # s1 = '12345' s2 = '45678' difference(s1, s2) = '123'
    def difference(self, s1, s2):
        s1_size = len(s1)
        s2_size = len(s2)
        for i in range(min(s1_size, s2_size), 0, -1):
            if s1[s1_size - i :] == s2[:i]:
                return s1[: s1_size - i]
        return s1

    def compose_doc(self, docs: List[Document], num) -> str:
        chunks = f"=== {docs[0].metadata[self.title_key]} ===\n\n"
        for doc in docs:
            chunks = self.difference(chunks, doc.page_content) + doc.page_content
        return f"NUMBER:{num}\n" + chunks

    def document(self, docs: List[Document]) -> str:
        # group by link
        docs.sort(key=lambda d: d.metadata[self.link_key])
        groups = itertools.groupby(docs, key=lambda d: d.metadata[self.link_key])
        # sort by idx and compose doc
        doc_str = ""
        i = 0
        for _, g in groups:
            i += 1
            doc_str += (
                self.compose_doc(sorted(g, key=lambda d: d.metadata[self.idx_key]), i)
                + "\n"
            )
        return doc_str