import subprocess
import webbrowser
from typing import Dict, Any
from flows.base_flows.atomic import AtomicFlow
import os

class CodeFileEditorAtomicFlow(AtomicFlow):
    def _write_code_to_temp_file(self, code_str, language_of_code, code_lib_location):
        directory = os.path.dirname(code_lib_location)
        if language_of_code.lower() == 'python':
            temp_file_location = os.path.join(directory, 'temp.py')
        else:
            raise NotImplemented
        try:
            content = (
                    "The below code will be appended to " +
                    code_lib_location + "\n"
                    "Edit the code directly or provide your thoughts down below if you have any suggestions.\n"
                    "If you prefer not to give any thoughts, just leave it blank, the code will be added to the library as it is "
                    "in this file.\n"
                    "###########\n"
                    "Code:\n" +
                    code_str +
                    "\n############\n"
                    "Thoughts:"
            )
            with open(temp_file_location, "w") as file:
                file.write(content)
            return True, f"Code written to {temp_file_location}", temp_file_location

        except Exception as e:
            return False, str(e), temp_file_location

    def _check_input(self, input_data: Dict[str, Any]):
        assert "code" in input_data, "code is not passed to CodeFileEditAtomicFlow"
        assert "language_of_code" in input_data, "language_of_code is not passed to CodeFileEditAtomicFlow"
        assert "memory_files" in input_data, "memory_files is not passed to CodeFileEditAtomicFlow"
        assert "code_library" in input_data["memory_files"], "code_library not in memory files"
        code_lib_loc = input_data["memory_files"]["code_library"]
        assert os.path.exists(code_lib_loc), f"{code_lib_loc} does not exist"
        assert os.path.isfile(code_lib_loc), f"{code_lib_loc} is not a file"

    def run(
            self,
            input_data: Dict[str, Any]
    ):
        self._check_input(input_data)
        code_str = input_data['code']
        language_of_code = input_data["language_of_code"]
        code_lib_location = input_data["memory_files"]["code_library"]
        result, code_editor_output, temp_file_location = self._write_code_to_temp_file(code_str, language_of_code, code_lib_location)
        if result:
            try:
                subprocess.run(["code", temp_file_location], timeout=10)
            except (subprocess.CalledProcessError, subprocess.TimeoutExpired):
                webbrowser.open(temp_file_location)
        response = {}
        response["code_editor_output"] = code_editor_output
        response["temp_code_file_location"] = temp_file_location
        return response