from typing import Dict, Any

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


class CtrlExMem_JarvisFlow(CtrlExMemFlow):
    def _on_reach_max_round(self):
        self._state_update_dict({
            "result": "the maximum amount of rounds was reached before the Jarvis flow has done the job",
            "summary": "JarvisFlow: 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": "Jarvis: " + output_payload["command_args"]["summary"],
                "status": "finished"
            }
        elif command == "manual_finish":
            # ~~~ return the manual quit status ~~~
            return {
                "EARLY_EXIT": True,
                "result": "JarvisFlow was terminated explicitly by the user, process is unfinished",
                "summary": "Jarvis: 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

        else:
            return output_payload