from dataclasses import dataclass, field
from typing import Any, Dict, Optional

from just_agents.llm_session import LLMSession
from just_agents.utils import load_config


@dataclass
class ChatAgent(LLMSession):
    """
    Default agent class, it assumes that each agent has a role, goal, and background story.
    """
    role: str  # "Role of the agent"
    goal: Optional[str] = None  # "Goal of the agent"
    task: Optional[str] = None # Tasks of the agent
    backstory: Optional[str] = None  # "Backstory of the agent"
    use_mako: bool = True #use mako for templating
    config: Dict[str, Any] = field(default_factory=lambda: load_config("agent_prompts.yaml"))  # "Configuration dictionary for the agent, usually loaded from yaml"
    agent_prompt: Optional[str] = None  # "Prompt used by the agent, usually generated by the template"

    def _render_template(self, template: str, mako: bool = False):
        """
        renders agent template
        :param template:
        :param mako:
        :return:
        """
        if mako:
            from mako.template import Template
            temp = Template(template)
            return temp.render(role = self.role,
                               goal = self.goal,
                               backstory = self.backstory,
                               task = self.task)
        else:
            return template.format(role = self.role,
                                  goal = self.goal,
                                  backstory = self.backstory,
                                  task = self.task
                                  )

    def __post_init__(self):
        """
        Post init logic
        :return:
        """
        super().__post_init__()
        agent_template: str = str(self.config["agent_prompt"])
        self.agent_prompt = self._render_template(agent_template, self.use_mako)
        if self.agent_prompt is not None:
            self.instruct(self.agent_prompt)