import os

from colmind.agents.llm import get_llm
import colmind.utils as U
from colmind.prompts import load_prompt

from langchain.schema import SystemMessage, HumanMessage

SYSTEM_TEMPLATE_PATH = ""
RESPONSE_TEMPLATE_PATH = ""

class EpisodicManager:
    def __init__(
        self,
        model_name = "gpt4",
        temperature = 0,
        retrieval_last_k = 5,
        request_timeout = 120,
        ckpt_dir = "episodic",
        username = "",
    ):

        self.llm = get_llm(model_name)
        # create folder for episodic data
        U.f_mkdir(f"{ckpt_dir}/episodes/")
        self.temperature = temperature
        self.retrieval_last_k = retrieval_last_k
        self.request_timeout = request_timeout
        self.ckpt_dir = ckpt_dir
        self.username = username
        self.episode_count = 0

    def get_episode(self, episode_id):
        # make the json to string
        return U.f_read(f"{self.ckpt_dir}/episodes/{episode_id}.txt")

    @property
    def episodes(self):
        episodes = ""
        for idx in range(self.episode_count - self.retrieval_last_k, self.episode_count):
            episodes += f"Episode {idx}:\n\n"
            episodes += self.get_episode(idx)
            episodes += "\n\n"

        return episodes

    def add_episode(self, episode):
        self.episode_count += 1
        U.f_write(f"{self.ckpt_dir}/episodes/{self.episode_count}.txt", episode)

    def render_system_message(self):

        system_template = load_prompt(SYSTEM_TEMPLATE_PATH)
        response_format = load_prompt(RESPONSE_TEMPLATE_PATH)

        system_message_prompt = SystemMessagePromptTemplate.from_template(
            system_template
        )
        system_message = system_message_prompt.format(
            episodes=self.episodes,
            response_format = response_format
        )

        assert isinstance(system_message, SystemMessage)
        return system_message

    def render_human_message(self):
        human_message = HumanMessage(
            message="Provide the list of insights below:",
        )

        assert isinstance(human_message, HumanMessage)
        return human_message

    def get_episodic_insights(self):
        system_message = self.render_system_message()
        human_message = self.render_human_message()

        response = self.llm([
            system_message,
            human_message
        ])
        return response.content
