from langchain import LLMChain, PromptTemplate
from langchain.chains import (
    ConversationalRetrievalChain,
    ConversationChain,
    RetrievalQA,
)
from langchain.chains.base import Chain
from langchain.memory import ConversationBufferMemory
from langchain.schema import BaseRetriever

from edu_assistant.learning_tasks.base import BaseTask
from edu_assistant.utils.langchain_utils import load_gpt4_llm, load_llm

TEMPLATE_CHAT = """The following is a friendly conversation between a human and an ai.
The ai is talkative and provides lots of specific details from its context.
If the ai does not know the answer to a question, it truthfully says it does not know.
The ai act following below instructions:
---
{instruction}
---


Current conversation:
{{chat_history}}
Human: {{input}}
AI:"""

TEMPLATE_CHAT_CONTEXT = """The following is a friendly conversation between a human and an ai.
The ai is talkative and provides lots of specific details from its context.
If the ai does not know the answer to a question, it truthfully says it does not know.
The ai act following below instructions:
---
{instruction}
---
Useful context for you to answer the question:
---
{{context}}
---

Current conversation:
{{chat_history}}
Human: {{question}}
AI:"""

TEMPLATE_ONCE = """The following is a friendly conversation between a human and an ai.
The ai is talkative and provides lots of specific details from its context.
If the ai does not know the answer to a question, it truthfully says it does not know.
The ai act following below instructions:
---
{instruction}
---

{{input}}
"""

TEMPLATE_ONCE_CONTEXT = """The following is a friendly conversation between a human and an ai.
The ai is talkative and provides lots of specific details from its context.
If the ai does not know the answer to a question, it truthfully says it does not know.
The ai act following below instructions:
---
{instruction}
---
Useful context for you to answer the question:
---
{{context}}
---

{{input}}
"""

DEFAULT_INSTRUCTION = """Act as a c++ professional to answer student aged 5-10 questions. Answer properly and politely.
Don't extend conversation multiple times. Only add one time saying."""


class QaTask(BaseTask):
    _session_store: dict
    _knowledge: BaseRetriever | None
    _qa_once: Chain

    HISTORY_KEY = "chat_history"

    def __init__(self, instruction: str = DEFAULT_INSTRUCTION, knowledge: BaseRetriever = None, enable_gpt4=False):
        """Create a new QaTask service.

        Args:
            instruction (str, optional): Instruction for this task ai. Defaults to "".
            knowledge (BaseRetriever, optional): Answer question with this knowledge retriever.
                If not set, will not use knowledge to answer question.
                Defaults to None.
            session_store (dict, optional): chat history store. Defaults to None.
                If not set, will use internal memory to store chat history. Which will be lost after restart and might
                cost huge memory.
        """
        self.enable_gpt4 = enable_gpt4
        # TODO: load threshold key from implement. value from config
        self.vectordbkwargs = {"score_threshold": 0.9}  # Qdrant cosine. higher is better.

        if knowledge:
            self._chat_prompt = PromptTemplate.from_template(TEMPLATE_CHAT_CONTEXT.format(instruction=instruction))
            self._once_prompt = PromptTemplate.from_template(TEMPLATE_ONCE_CONTEXT.format(instruction=instruction))
            self._input_key = "question"
            self._output_key = "answer"
        else:
            self._chat_prompt = PromptTemplate.from_template(TEMPLATE_CHAT.format(instruction=instruction))
            self._once_prompt = PromptTemplate.from_template(TEMPLATE_ONCE.format(instruction=instruction))
            self._input_key = "input"
            self._output_key = "response"

        self._session_store = {}
        self._knowledge = knowledge

        self._init_llm()

        self._qa_once = self._build_once_chain()

    def ask(
        self,
        question: str,
        session: bool = True,
        session_id: str = None,
        session_mem: ConversationBufferMemory | None = None,
    ) -> dict:
        """ask a question with chat history.

        Args:
            question (str): question to llm.
            session (bool, optional): whether use and store chat history. Defaults to False.
                if session_id is not set, a new session will be created.
            session_id (str, optional): specify a history qa session. Defaults to None.
            session_mem (list | None, optional): specify session memory. Defaults to None.
                if session_id is also set, memory will be replaced by passed one.

        Returns:
            dict: question answer and metadata.
                contains answer.
                contains session_id if session is True.
        """

        if session:
            args = {self._input_key: question}
            if session_id and session_id in self._session_store:
                chain = self._session_store[session_id]
            else:
                session_id = self._create_session_id()
                chain = self._create_session_chain(session_id)
        else:
            args = {self._input_key: question, QaTask.HISTORY_KEY: ""}
            chain = self._qa_once

        if session_mem:
            chain.memory = session_mem

        # TODO: ConversationalRetrievalChain should support vectordbkwargs
        # if self._knowledge:
        #     args["vectordbkwargs"] = self.vectordbkwargs
        result = chain(args)

        if session_id:
            result["session_id"] = session_id

        return result

    def _init_llm(self):
        self._main_llm = load_gpt4_llm() if self.enable_gpt4 else load_llm()
        self._secondary_llm = load_llm()

    def _build_once_chain(self):
        if not self._knowledge:
            return LLMChain(
                llm=self._main_llm,
                prompt=self._once_prompt,
            )
        else:
            return RetrievalQA.from_llm(
                llm=self._main_llm,
                retriever=self._knowledge,
                return_source_documents=True,
                prompt=self._once_prompt,
            )

    def _build_chat_chain(self):
        if not self._knowledge:
            return ConversationChain(
                llm=self._main_llm,
                memory=ConversationBufferMemory(
                    memory_key=QaTask.HISTORY_KEY, output_key=self._output_key, return_messages=True
                ),
                prompt=self._chat_prompt,
            )
        else:
            return ConversationalRetrievalChain.from_llm(
                llm=self._main_llm,
                retriever=self._knowledge,
                condense_question_llm=self._secondary_llm,
                return_source_documents=True,
                combine_docs_chain_kwargs={"prompt": self._chat_prompt},
                memory=ConversationBufferMemory(
                    memory_key=QaTask.HISTORY_KEY, output_key=self._output_key, return_messages=True
                ),
            )

    def _create_session_chain(self, session_id) -> ConversationChain:
        chain = self._build_chat_chain()
        self._session_store[session_id] = chain
        return chain