# -*- coding: utf-8 -*-
"""An agent that replies in a dictionary format."""
from typing import Optional

from ..message import Msg
from .agent import AgentBase
from ..parsers import ParserBase


class DictDialogAgent(AgentBase):
    """An agent that generates response in a dict format, where user can
    specify the required fields in the response via specifying the parser

    About parser, please refer to our
    [tutorial](https://modelscope.github.io/agentscope/en/tutorial/203-parser.html)

    For usage example, please refer to the example of werewolf in
    `examples/game_werewolf`"""

    def __init__(
        self,
        name: str,
        sys_prompt: str,
        model_config_name: str,
        use_memory: bool = True,
        memory_config: Optional[dict] = None,
        max_retries: Optional[int] = 3,
    ) -> None:
        """Initialize the dict dialog agent.

        Arguments:
            name (`str`):
                The name of the agent.
            sys_prompt (`Optional[str]`, defaults to `None`):
                The system prompt of the agent, which can be passed by args
                or hard-coded in the agent.
            model_config_name (`str`, defaults to None):
                The name of the model config, which is used to load model from
                configuration.
            use_memory (`bool`, defaults to `True`):
                Whether the agent has memory.
            memory_config (`Optional[dict]`, defaults to `None`):
                The config of memory.
            max_retries (`Optional[int]`, defaults to `None`):
                The maximum number of retries when failed to parse the model
                output.
        """  # noqa
        super().__init__(
            name=name,
            sys_prompt=sys_prompt,
            model_config_name=model_config_name,
            use_memory=use_memory,
            memory_config=memory_config,
        )

        self.parser = None
        self.max_retries = max_retries

    def set_parser(self, parser: ParserBase) -> None:
        """Set response parser, which will provide 1) format instruction; 2)
        response parsing; 3) filtering fields when returning message, storing
        message in memory. So developers only need to change the
        parser, and the agent will work as expected.
        """
        self.parser = parser

    def reply(self, x: dict = None) -> dict:
        """Reply function of the agent.
        Processes the input data, generates a prompt using the current
        dialogue memory and system prompt, and invokes the language
        model to produce a response. The response is then formatted
        and added to the dialogue memory.

        Args:
            x (`dict`, defaults to `None`):
                A dictionary representing the user's input to the agent.
                This input is added to the dialogue memory if provided.
        Returns:
            A dictionary representing the message generated by the agent in
            response to the user's input. It contains at least a 'speak' key
            with the textual response and may include other keys such as
            'agreement' if provided by the language model.

        Raises:
            `json.decoder.JSONDecodeError`:
                If the response from the language model is not valid JSON,
                it defaults to treating the response as plain text.
        """
        # record the input if needed
        if self.memory:
            self.memory.add(x)

        # prepare prompt
        prompt = self.model.format(
            Msg("system", self.sys_prompt, role="system"),
            self.memory
            and self.memory.get_memory()
            or x,  # type: ignore[arg-type]
            Msg("system", self.parser.format_instruction, "system"),
        )

        # call llm
        res = self.model(
            prompt,
            parse_func=self.parser.parse,
            max_retries=self.max_retries,
        )

        # Filter the parsed response by keys for storing in memory, returning
        # in the reply function, and feeding into the metadata field in the
        # returned message object.
        self.memory.add(
            Msg(self.name, self.parser.to_memory(res.parsed), "assistant"),
        )

        msg = Msg(
            self.name,
            content=self.parser.to_content(res.parsed),
            role="assistant",
            metadata=self.parser.to_metadata(res.parsed),
        )
        self.speak(msg)

        return msg
