from typing import Any, Dict, List
from flow_modules.Tachi67.ExtendLibraryFlowModule import ControllerFlow_ExtLib


from dataclasses import dataclass


@dataclass
class Command:
    name: str
    description: str
    input_args: List[str]

class Controller_CoderFlow(ControllerFlow_ExtLib):
    def __init__(
            self,
            commands: List[Command],
            **kwargs):
        super().__init__(**kwargs)
        self.system_message_prompt_template = self.system_message_prompt_template.partial(
            commands=self.build_commands_manual(commands),
            plan="no plans yet",
            plan_file_location="no location yet",
            code_library="no code library yet",
            code_library_location="no location yet",
            logs="no logs yet",
        )
        self.hint_for_model = """
        Make sure your response is in the following format:
              Response Format:
              {
              "command": "call one of the subordinates",
              "command_args": {
                  "arg name": "value"
                  }
              }
        """

    def _get_code_library_location(self, input_data):
        assert "memory_files" in input_data, "memory_files not passed to Coder/Controller"
        assert "code_library" in input_data["memory_files"], "code_library not in memory files"
        return input_data["memory_files"]["code_library"]

    def _get_code_library_content(self, input_data):
        assert "code_library" in input_data, "code_library not passed to Coder/Controller"
        code_library_content = input_data["code_library"]
        if len(code_library_content) == 0:
            code_library_content = 'No plan yet'
        return code_library_content

    def _update_prompts_and_input(self, input_data: Dict[str, Any]):
        if 'goal' in input_data:
            input_data['goal'] += self.hint_for_model
        if 'result' in input_data:
            input_data['result'] += self.hint_for_model
        plan_file_location = self._get_plan_file_location(input_data)
        plan_content = self._get_plan_content(input_data)
        code_library_location = self._get_code_library_location(input_data)
        code_library_content = self._get_code_library_content(input_data)
        logs_content = self._get_logs_content(input_data)
        self.system_message_prompt_template = self.system_message_prompt_template.partial(
            plan_file_location=plan_file_location,
            plan=plan_content,
            code_library_location=code_library_location,
            code_library=code_library_content,
            logs=logs_content
        )