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

from flow_modules.aiflows.AbstractBossFlowModule import AbstractBossFlow
class ExtendLibraryFlow(AbstractBossFlow):
    """ExtendLibraryFlow is one branch executor of CoderFlow. At a higher level, it takes the goal
    from the Coder, writes functions in an interactive way, test the code and append the newly written function to
    the code library.

    Workflow of ExtendLibrary:
    0. Coder calls ExtendLibrary with a goal.
    1. MemoryReading reads logs, plan, and code library.
    2. Planner makes a plan based on the goal.
    3. Write code in an interactive fashion.
    4. Test code.
    5. Finish, writes code to the library.
    Step 3-5 is done via prompting the controller.

    *Input Interface (expected input)*
    - `goal` (str): The goal from the caller (source flow, i.e. CoderFlow)
    *Output Interface (expected output)*
    - `result` (str): The result of the flow, the result will be returned to the caller (i.e. CoderFlow).
    - `summary` (str): The summary of the flow, the summary will be logged into the logs of the caller flow (i.e. CoderFlow).
    """
    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