from __future__ import annotations

import json
from datetime import datetime

from langchain_openai.chat_models import ChatOpenAI

from synthetic_agents.common.config import MESSAGE_DATETIME_FORMAT, settings
from synthetic_agents.model.constants import (
    DEFAULT_LLM_TEMPERATURE,
    DEFAULT_LLM_TOP_P,
    DEFAULT_OPENAI_CHAT_MODEL_NAME,
)
from synthetic_agents.model.entity.life_memory import LifeMemory
from synthetic_agents.prompt.builder import PromptBuilder
from synthetic_agents.prompt.loader import LocalAssetPromptTemplateLoader


class LifeFactGenerator:
    """
    This class represents a generator that creates random life memories and associated
    emotional states for an agent.
    """

    def __init__(
        self,
        user_profile: dict[str, str],
        reference_timestamp: datetime = datetime.now(),
        llm_name: str = DEFAULT_OPENAI_CHAT_MODEL_NAME,
        temperature: int = DEFAULT_LLM_TEMPERATURE,
        top_p: float = DEFAULT_LLM_TOP_P,
        api_key: str = settings.open_ai_api_key,
    ):
        """
        Creates a list of timestamped life memories with associated emotional states.

        :param user_profile: agent attributes in a json format to generate consistent life
            memories.
        :param reference_timestamp: a timestamp to use as reference. For instance, if the agent is
            20 years old and the reference_timestamp is 2020-10-01, it means the agent was born in
            2000 and memories before that should not exist.
        :param llm_name: name of the LLM model used to generate the memories. Defaults to
            DEFAULT_OPENAI_CHAT_MODEL_NAME.
        :param temperature: a non-negative value that controls the level of randomness of the
            responses generated by the LLM. The higher, the more random the responses. Defaults to
            DEFAULT_LLM_TEMPERATURE.
        :param top_p: the top-p next tokens to consider for sampling. Defaults to
            DEFAULT_LLM_TOP_P.
        :param api_key: key for authentication purposes when using the language model's API.
            Defaults to settings.open_ai_api_key.
        :raise ValueError: if percentages do not sum up to one.
        :return: a list of life memories.
        """
        self.user_profile = user_profile
        self.reference_timestamp = reference_timestamp
        self.llm_name = llm_name
        self.temperature = temperature
        self.top_p = top_p
        self.api_key = api_key

    def generate(self, num_memories: int) -> list[LifeMemory]:
        """
        Generates a list of random timestamped life facts with associated emotional states.

        :param num_memories: number of life facts to sample.
        :raise ValueError: if the LLM produces a malformed json object that cannot be converted to
            a life memory.
        :return: a list of randomly generated life facts.
        """
        if num_memories == 0:
            return []

        llm = ChatOpenAI(
            openai_api_key=self.api_key,
            model=self.llm_name,
            temperature=self.temperature,
            model_kwargs=dict(top_p=self.top_p),
        )

        prompt_builder = PromptBuilder(
            prompt_template_loader=LocalAssetPromptTemplateLoader("prompt/life_facts.txt"),
            placeholder_values=dict(
                num_memories=num_memories,
                user_profile="\n".join([f"{k}: {v}" for k, v in self.user_profile.items()]),
                reference_timestamp=self.reference_timestamp,
                timestamp_format=MESSAGE_DATETIME_FORMAT,
            ),
        )
        prompt = prompt_builder.build_prompt()
        response = llm.invoke(prompt)
        if not isinstance(response, str):
            response = response.content

        life_facts = [LifeMemory.from_json(life_fact) for life_fact in json.loads(response)]

        return life_facts
