import json
from copy import deepcopy
from typing import Any, Dict
from flow_modules.aiflows.ChatFlowModule import ChatAtomicFlow

class PlanGeneratorAtomicFlow(ChatAtomicFlow):
    """This class wraps around the Chat API to generate plan from a goal.

    *Input Interface Non Initialized*:
    - `goal`

    *Input Interface Initialized*:
    - `goal`

    *Output Interface*:
    - `plan`

    *Configuration Parameters*:
    - Also refer to ChatAtomicFlow (https://huggingface.co/aiflows/ChatFlowModule/blob/main/ChatAtomicFlow.py)
    - `input_interface_non_initialized`: The input interface when the conversation is not initialized.
    - `input_interface_initialized`: The input interface when the conversation is initialized.
    - `output_interface`: The output interface.
    - `backend`: The backend to use for the Chat API.
    - `system_message_prompt_template`: The template for the system message prompt.
    - `human_message_prompt_template`: The template for the human message prompt.
    - `init_human_message_prompt_template`: The initial human message prompt.

    """
    def __init__(self, **kwargs):
        """
        This function instantiates the class.
        :param kwargs: The configuration parameters.
        :type kwargs: Dict[str, Any]
        """
        super().__init__(**kwargs)
        self.hint_for_model = """
        Make sure your response is in the following format:
              Response Format:
              {
              "plan": "A step-by-step plan to finish the given goal, each step of plan should contain full information about writing a function",
              }
        """

    @classmethod
    def instantiate_from_config(cls, config):
        """
        This function instantiates the class from a configuration.
        :param config: The configuration.
        :type config: Dict[str, Any]
        :return: The instantiated class.
        :rtype: ChatAtomicFlow
        """
        flow_config = deepcopy(config)

        kwargs = {"flow_config": flow_config}

        # ~~~ Set up prompts ~~~
        kwargs.update(cls._set_up_prompts(flow_config))

        # ~~~ Set up backend ~~~
        kwargs.update(cls._set_up_backend(flow_config))

        # ~~~ Instantiate flow ~~~
        return cls(**kwargs)

    def _update_prompts_and_input(self, input_data: Dict[str, Any]):
        """
        This function updates the prompts and input data.
        :param input_data: The input data.
        :type input_data: Dict[str, Any]
        :return: None
        :rtype: None
        """
        if 'goal' in input_data:
            input_data['goal'] += self.hint_for_model

    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        This function runs the flow.
        :param input_data: The input data.
        :type input_data: Dict[str, Any]
        :return: The output data.
        :rtype: Dict[str, Any]
        """
        self._update_prompts_and_input(input_data)
        while True:
            api_output = super().run(input_data)["api_output"].strip()
            try:
                response = json.loads(api_output)
                return response
            except (json.decoder.JSONDecodeError, json.JSONDecodeError):
                new_goal = "The previous respond cannot be parsed with json.loads. Next time, do not provide any comments or code blocks. Make sure your next response is purely json parsable."
                new_input_data = input_data.copy()
                new_input_data['goal'] = new_goal
                input_data = new_input_data