from typing import Dict, Any

from flow_modules.Tachi67.AbstractBossFlowModule import CtrlExMemFlow
from aiflows.base_flows import CircularFlow

class CtrlExMem_CoderFlow(CtrlExMemFlow):
    """This class inherits from the CtrlExMemFlow class from AbstractBossFlowModule.
    See: https://huggingface.co/Tachi67/AbstractBossFlowModule/blob/main/CtrlExMemFlow.py

    *Input Interface*:
    - `plan`
    - `logs`
    - `code_library`: the signatures and docstring of the functions in the code library.
    - `memory_files`
    - `goal`

    *Output Interface*
    - `result` (str): The result of the flow, the result will be returned to the caller.
    - `summary` (str): The summary of the flow, the summary will be logged into the logs of the caller flow.
    
    """
    def _on_reach_max_round(self):
        self._state_update_dict({
            "result": "the maximum amount of rounds was reached before the coder flow has done the job",
            "summary": "CoderFlow: the maximum amount of rounds was reached before the flow has done the job",
            "status": "unfinished"
        })

    @CircularFlow.output_msg_payload_processor
    def detect_finish_or_continue(self, output_payload: Dict[str, Any], src_flow) -> Dict[str, Any]:
        command = output_payload["command"]
        if command == "finish":
            return {
                "EARLY_EXIT": True,
                "result": output_payload["command_args"]["summary"],
                "summary": "Coder: " + output_payload["command_args"]["summary"],
                "status": "finished"
            }
        elif command == "manual_finish":
            # ~~~ return the manual quit status ~~~
            return {
                "EARLY_EXIT": True,
                "result": "CoderFlow was terminated explicitly by the user, process is unfinished",
                "summary": "Coder: process terminated by the user explicitly, nothing generated",
                "status": "unfinished"
            }
        elif command == "update_plan":
            keys_to_fetch_from_state = ["memory_files"]
            fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
            output_payload["command_args"]["memory_files"] = fetched_state["memory_files"]
            return output_payload

        elif command == "re_plan":
            keys_to_fetch_from_state = ["plan", "memory_files"]
            fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
            output_payload["command_args"]["plan_file_location"] = fetched_state["memory_files"]["plan"]
            output_payload["command_args"]["plan"] = fetched_state["plan"]
            return output_payload
        
        elif command == "run_code":
            keys_to_fetch_from_state = ["memory_files"]
            fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
            output_payload["command_args"]["memory_files"] = fetched_state["memory_files"]
            return output_payload

        else:
            return output_payload