import abc
import os
import time
import urllib
from queue import Queue
from threading import Thread

from langchain.callbacks.tracers import LangChainTracer
from langchain.chains.base import Chain

from app_modules.llm_loader import LLMLoader, TextIteratorStreamer
from app_modules.utils import remove_extra_spaces


class LLMInference(metaclass=abc.ABCMeta):
    llm_loader: LLMLoader
    chain: Chain

    def __init__(self, llm_loader):
        self.llm_loader = llm_loader
        self.chain = None

    @abc.abstractmethod
    def create_chain(self) -> Chain:
        pass

    def get_chain(self, tracing: bool = False) -> Chain:
        if self.chain is None:
            if tracing:
                tracer = LangChainTracer()
                tracer.load_default_session()

            self.chain = self.create_chain()

        return self.chain

    def call_chain(
        self, inputs, streaming_handler, q: Queue = None, tracing: bool = False
    ):
        print(inputs)

        if self.llm_loader.streamer is not None and isinstance(
            self.llm_loader.streamer, TextIteratorStreamer
        ):
            self.llm_loader.streamer.reset(q)

        chain = self.get_chain(tracing)
        result = (
            self._run_qa_chain(
                chain,
                inputs,
                streaming_handler,
            )
            if streaming_handler is not None
            else chain(inputs)
        )

        if "answer" in result:
            result["answer"] = remove_extra_spaces(result["answer"])

            base_url = os.environ.get("PDF_FILE_BASE_URL")
            if base_url is not None and len(base_url) > 0:
                documents = result["source_documents"]
                for doc in documents:
                    source = doc.metadata["source"]
                    title = source.split("/")[-1]
                    doc.metadata["url"] = f"{base_url}{urllib.parse.quote(title)}"

        return result

    def _run_qa_chain(self, qa, inputs, streaming_handler):
        que = Queue()

        t = Thread(
            target=lambda qa, inputs, q, sh: q.put(qa(inputs, callbacks=[sh])),
            args=(qa, inputs, que, streaming_handler),
        )
        t.start()

        if self.llm_loader.streamer is not None and isinstance(
            self.llm_loader.streamer, TextIteratorStreamer
        ):
            count = 2 if len(inputs.get("chat_history")) > 0 else 1

            while count > 0:
                try:
                    for token in self.llm_loader.streamer:
                        streaming_handler.on_llm_new_token(token)

                    self.llm_loader.streamer.reset()
                    count -= 1
                except Exception:
                    print("nothing generated yet - retry in 0.5s")
                    time.sleep(0.5)

        t.join()
        return que.get()