import json

from langchain_openai.chat_models import ChatOpenAI

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


class UserProfileGenerator:
    """
    This class represents a generator that creates a user profile in json format containing the
    following information.

    1. `name`: The person's full name.
    2. `age`: The person's age.
    3. `race`: The person's race.
    4. `gender`: The person's gender.
    5. `dob`: The person's birthdate in the format YYYY-MM-DD.
    6. `birthplace`: The location of the person's birth in the format Place X where X is a random
        number.
    7. `current_job`: The person's current occupation.
    8. `places_lived`: A comma-separated list of places where the person has lived in the format
        Place X, Place Y where X and Y are random numbers.
    9. `children`: Children the person has (e.g., 1 son and 1 daughter).
    10. `siblings`: Siblings the person has (e.g., 1 sister).
    11. `fathers_jobs`: A comma-separated list of jobs held by the person's father.
    12. `mothers_jobs`: A comma-separated list of jobs held by the person's mother.
    13. `father_adjectives`:  A comma-separated list of three adjectives describing the person's
        father. It can contain negative adjectives.
    14. `mother_adjectives`: A comma-separated list of three adjectives describing the person's
        mother. It can contain negative adjectives.
    """

    def __init__(
        self,
        llm_name: str = DEFAULT_OPENAI_CHAT_MODEL_NAME,
        temperature: float = DEFAULT_LLM_TEMPERATURE,
        top_p: float = DEFAULT_LLM_TOP_P,
        api_key: str = settings.open_ai_api_key,
    ):
        """
        Creates an agent attribute generator.

        :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.
        """
        self.llm_name = llm_name
        self.temperature = temperature
        self.top_p = top_p
        self.api_key = api_key

    def generate(self, num_profiles: int) -> list[dict[str, str]]:
        """
        Generates a list of profiles for different agents.

        Each profile is generated by a separate call to the LLM.

        :param num_profiles: number of profiles to generate.
        :return: a list of user profiles in json format.
        """

        if num_profiles == 0:
            return []

        prompt_builder = PromptBuilder(
            prompt_template_loader=LocalAssetPromptTemplateLoader("prompt/user_profile.txt")
        )
        llm = ChatOpenAI(
            openai_api_key=self.api_key,
            model=self.llm_name,
            temperature=self.temperature,
            model_kwargs=dict(top_p=self.top_p),
        )
        prompt = prompt_builder.build_prompt()

        user_profiles = []
        for _ in range(num_profiles):
            output = llm.invoke(prompt)
            if not isinstance(output, str):
                output = output.content
            response = json.loads(output)
            user_profiles.append(response)

        return user_profiles
