from typing import Dict, Any
import os
from flow_modules.aiflows.PlanWriterFlowModule import PlanWriterFlow
from aiflows.base_flows import CircularFlow



class Planner_CoderFlow(PlanWriterFlow):
    """Planner of the coder flow, it inherits from PlanWriterFlow, see: https://huggingface.co/aiflows/PlanWriterFlowModule
    This flow is responsible for generating the plan for the coder flow.

    *Input Interface*
    - `goal`
    - `memory_files`
    - `code_library`

    *Output Interface*
    - `plan`
    - `summary`
    - `status`

    *Configuration Parameters*:
    - Look at the configuration parameters of PlanWriterFlowModule for detailed info. (https://huggingface.co/aiflows/PlanWriterFlowModule)
    - `input_interface`: the input interface of the flow, default: `["goal", "memory_files", "code_library"]`
    - `output_interface`: the output interface of the flow, default: `["plan", "summary", "status"]`
    - `subflows_config`: configuration of the subflows.
    - `topology`: topology of the subflows.
    """
    def _on_reach_max_round(self):
        self._state_update_dict({
            "plan": "The maximum amount of rounds was reached before the model generated the plan.",
            "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":
            # ~~~ fetch temp file location, plan content, memory file (of upper level flow e.g. ExtLib) from flow state
            keys_to_fetch_from_state = ["temp_plan_file_location", "plan", "memory_files"]
            fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
            temp_plan_file_location = fetched_state["temp_plan_file_location"]
            plan_content = fetched_state["plan"]
            plan_file_location = fetched_state["memory_files"]["plan"]

            # ~~~ delete the temp plan file ~~~
            if os.path.exists(temp_plan_file_location):
                os.remove(temp_plan_file_location)

            # ~~~ write plan content to plan file ~~~
            with open(plan_file_location, 'w') as file:
                file.write(plan_content)

            # ~~~ return the plan content ~~~
            return {
                "EARLY_EXIT": True,
                "plan": plan_content,
                "summary": "Coder/PlanWriter: " + output_payload["command_args"]["summary"],
                "status": "finished"
            }
        elif command == "manual_finish":
            # ~~~ delete the temp plan file ~~~
            keys_to_fetch_from_state = ["temp_plan_file_location"]
            fetched_state = self._fetch_state_attributes_by_keys(keys=keys_to_fetch_from_state)
            temp_plan_file_location = fetched_state["temp_plan_file_location"]
            if os.path.exists(temp_plan_file_location):
                os.remove(temp_plan_file_location)
            # ~~~ return the manual quit status ~~~
            return {
                "EARLY_EXIT": True,
                "plan": "no plan was generated",
                "summary": "Coder/PlanWriter: PlanWriter was terminated explicitly by the user, process is unfinished",
                "status": "unfinished"
            }
        elif command == "write_plan":
            keys_to_fetch_from_state = ["memory_files", "code_library"]
            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"]["code_library_location"] = fetched_state["memory_files"]["code_library"]
            output_payload["command_args"]["code_library"] = fetched_state["code_library"]
            return output_payload
        else:
            return output_payload