import itertools
import re
from typing import List

import tiktoken
from langchain_core.documents import Document
from resolver import block

from blocks.base import BaseBlock
from patterns.embedding import Namespace

# from langchain_core.messages import AIMessage, HumanMessage
# from langchain_core.retrievers import BaseRetriever


@block(name="Overlap_Doc_Retrieve_Private", kind="tools")
class SlidingDocChainPrivate(BaseBlock):
    def __init__(
        self,
        retriever: Namespace,
        max_tokens: int = 1024,
        link_key: str = "link",
        title_key: str = "title",
        idx_key: str = "idx",
        text_key: str = "content",
    ):
        super(SlidingDocChainPrivate, 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
        self._text_key = text_key

    def __call__(self, text: str) -> str:
        retrieved_chunks = self.retriever.retrieve(text, limit=16)
        for chunk in retrieved_chunks:
            chunk["namespace"] = self.retriever._name
        relevant = sorted(retrieved_chunks, key=lambda x: x["_score"], reverse=True)[
            :16
        ]
        docs = []
        for r in relevant:
            metadata = {
                "content": r["content"],
                "title": r["title"],
                "link": r["link"],
                "idx": r["idx"],
                "namespace": r["namespace"],
            }
            if self._text_key in metadata:
                text = metadata.pop(self._text_key)
                text = self.remove_images_from_markdown(text)
                text = self.remove_links_from_markdown(text)
                docs.append(Document(page_content=text, metadata=metadata))
        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

    def remove_images_from_markdown(self, markdown_text):
        pattern = r"!\[.*?\]\(.*?\)"
        replaced_text = re.sub(pattern, "", markdown_text)
        return replaced_text

    def remove_links_from_markdown(self, markdown_text):
        pattern = r"\[(.*?)\]\(.*?\)"
        replaced_text = re.sub(pattern, r"\1", markdown_text)
        return replaced_text