""" Module description

Attributes:
    IS_SUBQUESTION_PREDICATE (str):
    default_prompt (PromptTemplate):

"""

from typing import Optional
from pathlib import Path
import time

from langchain_core.language_models import BaseLanguageModel
from langchain_core.runnables import (
    RunnablePassthrough,
    RunnableLambda,
)
from langchain_core.tools import Tool
from langchain_core.prompts.base import BasePromptTemplate
from langchain_core.prompts import PromptTemplate

from langchain_core.pydantic_v1 import BaseModel, Field

from motleycrew.tools import MotleyTool
from motleycrew.common import LLMFramework
from motleycrew.common.llms import init_llm
from motleycrew.common.utils import print_passthrough
from motleycrew.storage import MotleyGraphStore
from motleycrew.common import logger


from motleycrew.applications.research_agent.question import Question, QuestionGenerationTaskUnit

IS_SUBQUESTION_PREDICATE = "is_subquestion"

default_prompt = PromptTemplate.from_template(
    """
You are a part of a team. The ultimate goal of your team is to
answer the following Question: '{question}'.\n
Your team has discovered some new text (delimited by ```) that may be relevant to your ultimate goal.
text: \n ``` {context} ``` \n
Your task is to ask new questions that may help your team achieve the ultimate goal.
If you think that the text is relevant to your ultimate goal, then ask new questions.
New questions should be based only on the text and the goal Question and no other previous knowledge.

You can ask up to {num_questions} new questions.
Return the questions each on a new line and ending with a single question mark.
Don't return anything else except these questions.
"""
)

# " The new questions should have no semantic overlap with questions in the following list:\n"
# " {previous_questions}\n"


class QuestionGeneratorTool(MotleyTool):
    """
    Gets a question as input
    Retrieves relevant docs (llama index basic RAG)
    (Retrieves existing questions from graph (to avoid overlap))
    Generates extra questions (research agent prompt)

    Adds questions as children of current q by calling Q insertion tool once
    exits
    """

    def __init__(
        self,
        query_tool: MotleyTool,
        graph: MotleyGraphStore,
        max_questions: int = 3,
        llm: Optional[BaseLanguageModel] = None,
        prompt: str | BasePromptTemplate = None,
    ):
        """Description

        Args:
            query_tool (MotleyTool):
            graph (MotleyGraphStore):
            max_questions (:obj:`int`, optional):
            llm (:obj:`BaseLanguageModel`, optional:
            prompt (:obj:`str`, :obj:`BasePromptTemplate`, optional):
        """
        langchain_tool = create_question_generator_langchain_tool(
            query_tool=query_tool,
            graph=graph,
            max_questions=max_questions,
            llm=llm,
            prompt=prompt,
        )

        super().__init__(langchain_tool)


class QuestionGeneratorToolInput(BaseModel, arbitrary_types_allowed=True):
    """Input for the Question Generator Tool.

    Attributes:
        question (Question):
    """

    question: Question = Field(description="The input question for which to generate subquestions.")


def create_question_generator_langchain_tool(
    query_tool: MotleyTool,
    graph: MotleyGraphStore,
    max_questions: int = 3,
    llm: Optional[BaseLanguageModel] = None,
    prompt: str | BasePromptTemplate = None,
):
    """Description

    Args:
        query_tool (MotleyTool):
        graph (MotleyGraphStore):
        max_questions (:obj:`int`, optional):
        llm (:obj:`BaseLanguageModel`, optional:
        prompt (:obj:`str`, :obj:`BasePromptTemplate`, optional):

    Returns:

    """
    if llm is None:
        llm = init_llm(llm_framework=LLMFramework.LANGCHAIN)

    llm.bind(json_mode=True)

    if prompt is None:
        prompt = default_prompt
    elif isinstance(prompt, str):
        prompt = PromptTemplate.from_template(prompt)

    assert isinstance(prompt, BasePromptTemplate), "Prompt must be a string or a BasePromptTemplate"

    def insert_questions(input_dict) -> None:
        questions_raw = input_dict["subquestions"].content
        questions = [q.strip() for q in questions_raw.split("\n") if len(q.strip()) > 1]
        for q in questions:
            logger.info("Inserting question: %s", q)
            subquestion = graph.insert_node(Question(question=q))
            graph.create_relation(input_dict["question"], subquestion, IS_SUBQUESTION_PREDICATE)
        logger.info("Inserted %s questions", len(questions))

    def set_context(input_dict: dict):
        node = input_dict["question"]
        node.context = input_dict["context"]

    pipeline = (
        RunnableLambda(print_passthrough)
        | RunnablePassthrough().assign(context=query_tool.to_langchain_tool())
        | RunnableLambda(print_passthrough)
        | RunnablePassthrough().assign(
            subquestions=prompt.partial(num_questions=str(max_questions)) | llm
        )
        | RunnableLambda(print_passthrough)
        | {
            "set_context": RunnableLambda(set_context),
            "insert_questions": RunnableLambda(insert_questions),
        }
    )

    return Tool.from_function(
        func=lambda q: pipeline.invoke({"question": q}),
        name="Question Generator Tool",
        description="""Generate a list of questions based on the input question, 
    and insert them into the knowledge graph.""",
        args_schema=QuestionGeneratorToolInput,
    )


if __name__ == "__main__":
    import kuzu
    from llama_index.graph_stores.kuzu import KuzuGraphStore

    here = Path(__file__).parent
    db_path = str(here / "test2")

    db = kuzu.Database(db_path)
    graph_store = KuzuGraphStore(db)

    query_tool = MotleyTool.from_langchain_tool(
        Tool.from_function(
            func=lambda question: [
                "Germany has consisted of many different states over the years",
                "The capital of France has moved in 1815, from Lyons to Paris",
                "France actually has two capitals, one in the north and one in the south",
            ],
            name="Query Tool",
            description="Query the library for relevant information.",
            args_schema=QuestionGeneratorToolInput,
        )
    )

    tool = QuestionGeneratorTool(
        query_tool=query_tool,
        graph=graph_store,
        max_questions=3,
    )

    tool.invoke({"question": "What is the capital of France?"})
    print("Done!")
