import os
from typing import Dict, Any
from flow_modules.aiflows.CodeFileEditFlowModule import CodeFileEditAtomicFlow

class RunCodeFileEditAtomicFlow(CodeFileEditAtomicFlow):
    """Refer to: https://huggingface.co/aiflows/CodeFileEditFlowModule

    *Input Interface*:
    - `code`: str
    - `language_of_code`: str
    - `memory_files`: Dict[str, str]

    *Output Interface*:
    - `code_editor_output`: str
    - `temp_code_file_location`: str

    *Configuration Parameters*:
    - `input_interface`: The input interface of the atomic flow.
    - `output_interface`: The output interface of the atomic flow.

    """

    def _generate_content(self, code_str, comment_sign) -> str:
        """
        Generate the content to write to the temp file.
        :param code_str: The code string.
        :type code_str: str
        :param comment_sign: The comment sign for the language.
        :type comment_sign: str
        :return: The content to write to the temp file.
        :rtype: str
        """
        if comment_sign == "<!-- -->":
            instruction = (
                "<!-- The below code will be executed.\n"
                "Change the code if you consider necessary.\n"
                "Shut this vscode session completely to start running the code. -->\n"
            )
        else:
            instruction = (
                f"{comment_sign} The below code will be executed.\n"
                f"{comment_sign} Change the code if you consider necessary.\n"
                f"{comment_sign} Shut this vscode session completely to start running the code.\n"
            )
        content = (
                instruction +
                "###########\n"
                "# Code:\n" +
                code_str +
                "\n############\n"
        )
        return content

    def _generate_temp_file_location(self, code_lib_location, file_extension):
        """
        Generate the temp file location.
        :param code_lib_location: The code library location.
        :type code_lib_location: str
        :param file_extension: The file extension.
        :type file_extension: str
        :return: The temp file location.
        :rtype: str
        """
        directory = os.path.dirname(code_lib_location)
        ret = os.path.join(directory, 'run_code_temp'+file_extension)
        return ret

    def _check_input(self, input_data: Dict[str, Any]):
        """
        Check the input data.
        :param input_data: The input data.
        :type input_data: Dict[str, Any]
        :return: None
        :rtype: None
        :raises AssertionError: If the input data is not valid.
        """
        assert "code" in input_data, "code is not passed to RunCodeFileEditAtomicFlow"
        assert "memory_files" in input_data, "memory_files is not passed to RunCodeFileEditAtomicFlow"
        assert "language" in input_data, "language is not passed to RunCodeFileEditAtomicFlow"

    def _generate_file_extension_and_comment_sign(self, language: str):
        """
        Generate the file extension and comment sign for the language.
        :param language: The language.
        :type language: str
        :return: The file extension and comment sign.
        :rtype: Dict[str, str]
        :raises NotImplemented: If the language is not supported.
        """
        details = {
            "python": {"extension": ".py", "comment": "#"},
            "bash": {"extension": ".sh", "comment": "#"},
            "shell": {"extension": ".sh", "comment": "#"},
            "javascript": {"extension": ".js", "comment": "//"},
            "html": {"extension": ".html", "comment": "<!-- -->"},
            "applescript": {"extension": ".scpt", "comment": "--"},
            "r": {"extension": ".r", "comment": "#"},
            "powershell": {"extension": ".ps1", "comment": "#"}
        }
        if language.lower() not in details:
            raise NotImplemented
        return details.get(language.lower())

    def _generate_input_to_writer(self, input_data: Dict[str, Any]):
        """
        Generate the input to the writer.
        :param input_data: The input data.
        :type input_data: Dict[str, Any]
        :return: The input to the writer.
        :rtype: Tuple[str, str]
        :raises NotImplemented: If the language is not supported.
        """
        code_str = input_data['code']
        code_lib_location = input_data["memory_files"]["code_library"]
        
        # if we are running python code and if the code is importing the code library after one modification, we need
        # to re-import the library.
        code_lang = input_data["language"].lower()
        code_lib_name = os.path.splitext(os.path.basename(code_lib_location))[0]
        if code_lang == 'python' and code_lib_name in code_str and 'importlib.reload' not in code_str:
            prepend_code = f"import importlib\nimport {code_lib_name}\nimportlib.reload({code_lib_name})\n"
            code_str = prepend_code + code_str
        
        lang_details = self._generate_file_extension_and_comment_sign(input_data["language"])
        content_to_write = self._generate_content(code_str, lang_details["comment"])
        file_location_to_write = self._generate_temp_file_location(code_lib_location, lang_details["extension"])
        return content_to_write, file_location_to_write