from typing import Dict, Any
from aiflows.utils import logging
log = logging.get_logger(__name__)

from flow_modules.Tachi67.AbstractBossFlowModule import AbstractBossFlow

class CoderFlow(AbstractBossFlow):
    """Coder flow is one executor branch of the Jarvis flow. At a higher level, it is a flow that
    writes and runs code given a goal. In the Jarvis flow, the Coder flow in invoked by the controller,
    The Coder flow receives the goal generated by the controller, writes and runs code in an interactive fashion.

    The Coder flow has the similar structure as the Jarvis flow (inherited from AbstractBossFlow).


    *Input Interface (expected input)*
    - `goal` (str): The goal from the caller (source flow, i.e. JarvisFlow)
    *Output Interface (expected output)*
    - `result` (str): The result of the flow, the result will be returned to the caller (i.e. JarvisFlow).
    - `summary` (str): The summary of the flow, the summary will be logged into the logs of the caller flow (i.e. JarvisFlow).
    """
    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        # ~~~ sets the input_data in the flow_state dict ~~~
        self._state_update_dict(update_data=input_data)

        # ~~~ set the memory file to the flow state ~~~
        self._state_update_dict(update_data={"memory_files": self.memory_files})

        max_rounds = self.flow_config.get("max_rounds", 1)
        if max_rounds is None:
            log.info(f"Running {self.flow_config['name']} without `max_rounds` until the early exit condition is met.")

        self._sequential_run(max_rounds=max_rounds)

        output = self._get_output_from_state()

        self.reset(full_reset=True, recursive=True, src_flow=self)

        return output