import time
from copy import deepcopy
from typing import Annotated

import google.generativeai as genai
import instructor
from dotenv import load_dotenv
from pydantic import AfterValidator, BaseModel, Field

from burr.core import Application, ApplicationBuilder, State, expr
from burr.core.action import action
from burr.tracking import LocalTrackingClient

load_dotenv()  # load the GOOGLE_API_KEY from your .env file

ask_gemini = instructor.from_gemini(
    client=genai.GenerativeModel(model_name="gemini-1.5-flash-latest")
)

MIN_WORDS = 3
MAX_WORDS = 5
NUM_REQUIRED_TOPICS = 3
MIN_SUBTOPICS = 2
MAX_SUBTOPICS = 4
MIN_CONCEPTS = 2
MAX_CONCEPTS = 4
ATTEMPTS = 3
TEMPERATURE = 0.3


class Subtopic(BaseModel):
    name: Annotated[str, AfterValidator(str.title)]
    # pydantic will ensure that we have at least MIN_CONCEPTS and at most MAX_CONCEPTS concepts
    # if we don't, the LLM will be asked to generate the concepts again
    concepts: Annotated[
        list[str], AfterValidator(lambda concepts: [x.title() for x in concepts])
    ] = Field(
        description=f"{MIN_CONCEPTS}-{MAX_CONCEPTS} concepts covered in the subtopic.",
        min_length=MIN_CONCEPTS,
        max_length=MAX_CONCEPTS,
    )


class Topic(BaseModel):
    name: Annotated[str, AfterValidator(str.title)]
    # pydantic will ensure that we have at least MIN_SUBTOPICS and at most MAX_SUBTOPICS subtopics
    # if we don't, the LLM will be asked to generate the subtopics again
    subtopics: list[Subtopic] = Field(
        description=f"{MIN_SUBTOPICS}-{MAX_SUBTOPICS} ordered subtopics with concepts.",
        min_length=MIN_SUBTOPICS,
        max_length=MAX_SUBTOPICS,
    )

    def __str__(self) -> str:
        topic_str = f"TOPIC: {self.name}\n"
        for i, subtopic in enumerate(self.subtopics):
            topic_str += f"SUBTOPIC {i + 1}: {subtopic.name}\n"
            for j, concept in enumerate(subtopic.concepts):
                topic_str += f"CONCEPT {j + 1}: {concept}\n"
        return topic_str


def topics_system_template(
    num_required_topics: int = NUM_REQUIRED_TOPICS,
    min_words: int = MIN_WORDS,
    max_words: int = MAX_WORDS,
    min_subtopics: int = MIN_SUBTOPICS,
    max_subtopics: int = MAX_SUBTOPICS,
    min_concepts: int = MIN_CONCEPTS,
    max_concepts: int = MAX_CONCEPTS,
) -> str:
    return f"""\
"You are a world class course instructor."
You'll be given a course outline and you have to generate {num_required_topics} topics.
For each topic:
    1. Generate a {min_words}-{max_words} word topic name that encapsulates the description.
    2. Generate {min_subtopics}-{max_subtopics} subtopics for the topic. Also {min_words}-{max_words} words each.
    For each subtopic:
        Generate {min_concepts}-{max_concepts} concepts. Also {min_words}-{max_words} words each. The concepts should be related to the subtopic.
        Think of concepts as the smallest unit of knowledge that can be taught from the subtopic. And add a verb to the concept to make it actionable.
        For example:
            "Calculate Derivatives" instead of "Derivatives".
            "Identify Finite Sets" instead of "Finite Sets".
            "Find the y-intercept" instead of "y-intercept".
    The subtopics and concepts should be in the correct order.\
"""


def creation_template(
    outline: str,
    topics_so_far: list[Topic],
    num_required_topics: int = NUM_REQUIRED_TOPICS,
) -> str:
    prompt_strs = [
        f"<outline>\n{outline}\n</outline>",
        f"<num_required_topics>\n{num_required_topics}\n</num_required_topics>",
    ]
    topics_so_far_str = "<topics_so_far>\n"
    if topics_so_far:
        for i, topic in enumerate(topics_so_far):
            topics_so_far_str += f"{i + 1}/{num_required_topics}\n{topic}\n"
    topics_so_far_str = topics_so_far_str.strip() + "\n</topics_so_far>"
    prompt_strs.append(topics_so_far_str)
    prompt_strs.append("Generate the next topic.")
    return "\n\n".join(prompt_strs)


@action(reads=[], writes=["outline", "num_required_topics", "chat_history"])
def setup(state: State, outline: str, num_required_topics: int) -> State:
    """
    Initializes the state of the application with the course outline and the number of topics to be generated.
    """
    return state.update(
        outline=outline,
        num_required_topics=num_required_topics,
        chat_history=[
            {
                "role": "system",
                "content": topics_system_template(num_required_topics=num_required_topics),
            }
        ],
    )


@action(
    reads=[
        "chat_history",
        "outline",
        "num_required_topics",
        "topics_so_far",
        "topic_feedback",
    ],
    writes=["generated_topic", "chat_history"],
)
def creator(
    state: State,
    attempts: int = ATTEMPTS,
) -> tuple[dict, State]:
    """
    Generates a topic based on the outline and topics generated so far using the create function from the Instructor library.
    https://python.useinstructor.com/#using-gemini

    Parameters:
    - state (State): The current state of the application.
    - attempts (int, optional): The number of times the model should try to generate a response. 3 by default. If the model fails to generate a response after the specified number of attempts, the function will return with 'generated_topic' set to None.

    Returns:
    - tuple[dict, State]: A tuple containing the generated topic and the updated state.

    Instructor Parameters Used:
    - response_model: The Pydantic model or the type that we want the response to be. In our case, it's the `Topic` model.
    - messages: The chat history that we want to pass to the model. This could have the prompt for generating the next topic, or the previously generated topic with user feedback.
    - max_retries: The number of times we want to retry if the model fails to generate a response. 1 means just the original attempt, 2 means the original attempt and one retry.
    - temperature: The temperature of the model. A float between 0 and 1. Lower values give more deterministic results. 0.3 by default.
    """

    num_required_topics = state.get("num_required_topics", NUM_REQUIRED_TOPICS)
    topics_so_far = state.get("topics_so_far", [])
    topic_feedback = state.get("topic_feedback", "")
    messages = deepcopy(state["chat_history"])

    if topic_feedback:
        # this means that we had some user feedback on the previously generated topic
        # so instead of prompting for a new topic, we will add the feedback as a user message
        # the previously generated topic is already an assistant message at the end of the chat history
        user_message = f"User feedback: {topic_feedback.strip()}"
    else:
        # either no previous topic or no feedback on the previous topic
        # so we will prompt for a new topic based on the outline and topics generated so far
        user_message = creation_template(
            outline=state["outline"],
            topics_so_far=topics_so_far,
            num_required_topics=num_required_topics,
        )
    messages.append({"role": "user", "content": user_message})
    # these are the parameters that we will pass to the create function from the Instructor library
    create_kwargs = dict(
        response_model=Topic,
        messages=messages,
        max_retries=attempts,
        temperature=TEMPERATURE,
    )
    try:
        res = ask_gemini.create(**create_kwargs)  # type: ignore
    except Exception as e:
        print(f"ERROR in creator: {e}")
        # if the model fails to generate a response even after multiple attempts, we will return None
        # creator will be called again until a valid response is generated
        return {"generated_topic": None}, state.update(generated_topic=None)
    # add the user message and the generated topic to the chat history
    return {"generated_topic": res}, state.update(generated_topic=res).append(
        chat_history={"role": "user", "content": user_message}
    ).append(chat_history={"role": "assistant", "content": str(res)})


@action(
    reads=["generated_topic", "num_required_topics", "topics_so_far"],
    writes=["topic_feedback"],
)
def get_topic_feedback(state: State) -> tuple[dict, State]:
    """
    Asks the user for feedback on the generated topic.
    If the user is happy with the generated topic, they can leave the feedback empty.
    This feedback will be added as a user message to the chat history.
    """
    time.sleep(2)
    num_required_topics = state.get("num_required_topics", NUM_REQUIRED_TOPICS)
    num_topics_so_far = len(state.get("topics_so_far", []))
    generated_topic = state["generated_topic"]
    generated_topic_str = f"{num_topics_so_far + 1}/{num_required_topics}\n{generated_topic}"
    feedback = (
        input(
            f"Give your feedback on this generated topic. Leave empty if it's perfect.\n\n{generated_topic_str}"
        )
        or None
    )

    return {"topic_feedback": feedback}, state.update(topic_feedback=feedback)


@action(reads=["generated_topic"], writes=["topics_so_far"])
def update_topics_so_far(state: State) -> State:
    """
    This is called if there is no feedback on the generated topic.
    We assume that the user is happy with the generated topic and we add it to the list of topics generated so far.
    This list will be incorporated in the prompt for the next topic.
    """
    return state.append(topics_so_far=state["generated_topic"])


@action(reads=["topics_so_far"], writes=[])
def terminal(state: State) -> tuple[dict, State]:
    """
    This is the terminal state of the application.
    We will reach this state when the required number of topics have been generated.
    """
    return {"topics_so_far": state["topics_so_far"]}, state


def application(
    app_id: str | None = None,
    username: str | None = None,
    project: str = "topics_creator",
) -> Application:
    tracker = LocalTrackingClient(project=project)
    builder = (
        ApplicationBuilder()
        .with_actions(setup, creator, get_topic_feedback, update_topics_so_far, terminal)
        .with_transitions(
            ("setup", "creator"),
            # go to creator again if the generated topic is None
            ("creator", "creator", expr("generated_topic is None")),  # type: ignore
            # get feedback if the generated topic is not None
            ("creator", "get_topic_feedback"),
            # if the feedback is empty, add the generated topic to the list of topics generated so far
            (
                "get_topic_feedback",
                "update_topics_so_far",
                expr("topic_feedback is None"),  # type: ignore
            ),
            # if the feedback is not empty, go back to creator and add the feedback as a user message
            ("get_topic_feedback", "creator"),
            # if the number of topics generated so far is equal to the required number of topics, go to terminal
            (
                "update_topics_so_far",
                "terminal",
                expr("len(topics_so_far) == num_required_topics"),  # type: ignore
            ),
            # otherwise, go back to creator after updating the topics generated so far
            ("update_topics_so_far", "creator"),
        )
        .with_tracker("local", project=project)
        .with_identifiers(app_id=app_id, partition_key=username)  # type: ignore
        .initialize_from(
            tracker,
            resume_at_next_action=True,
            default_entrypoint="setup",
            default_state={},
        )
    )
    return builder.build()


if __name__ == "__main__":
    app = application()
    app.visualize(
        output_file_path="statemachine",
        include_conditions=True,
        include_state=False,
        format="png",
    )
