from copy import deepcopy
from typing import Any, Dict

from aiflows.base_flows import SequentialFlow
from aiflows.utils import logging

logging.set_verbosity_debug()
log = logging.get_logger(__name__)

class TestCodeFlow(SequentialFlow):
    """This class is used to test code. It is a sequential flow that runs the following steps:
    1. Prepares the code to be tested, it is composed of the code to be tested and necessary import statements manually added.
    2. Opens the code in VSCode and waits for the user to clode the vscode session. The user is able to add tests.
    3. The following will be tested:
        a. (Default & Compulsory) Code syntax;
        b. (Added by user) Any other tests.
    4. Runs the test and returns the output.
    
    *Input Interface*:
    - `code` (str): The code to be tested.
    
    *Output Interface*:
    - `feedback` (str): The test results.

    *Configuration Parameters*:
    - `memory_files`: The memory files to be used in the flow, typically we need to have the code library as memory file here.
    - `input_interface`: The input interface of the flow.
    - `output_interface`: The output interface of the flow.
    - `subflows_config`: The subflows configuration.
    - `topology`: The topology of the subflows.

    """
    REQUIRED_KEYS_CONFIG = ["max_rounds", "early_exit_key", "topology", "memory_files"]

    def __init__(
            self,
            memory_files: Dict[str, Any],
            **kwargs
    ):
        """
        This function instantiates the class.
        :param memory_files: The memory files to be used in the flow, typically we need to have the code library as memory file here.
        :type memory_files: Dict[str, Any]
        :param kwargs: The keyword arguments.
        :type kwargs: Any
        """
        super().__init__(**kwargs)
        self.memory_files = memory_files

    @classmethod
    def instantiate_from_config(cls, config):
        """
        This function instantiates the class from a configuration dictionary.
        :param config: The configuration dictionary.
        :type config: Dict[str, Any]
        :return: The instantiated class.
        :rtype: TestCodeFlow
        """
        flow_config = deepcopy(config)

        kwargs = {"flow_config": flow_config}

        # ~~~ Set up memory file ~~~
        memory_files = flow_config["memory_files"]
        kwargs.update({"memory_files": memory_files})

        # ~~~ Set up subflows ~~~
        kwargs.update({"subflows": cls._set_up_subflows(flow_config)})

        # ~~~ Instantiate flow ~~~
        return cls(**kwargs)

    def run(self, input_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        This function runs the flow.
        :param input_data: The input data.
        :type input_data: Dict[str, Any]
        :return: The output data.
        :rtype: Dict[str, Any]
        """
        # ~~~ sets the input_data in the flow_state dict ~~~
        self._state_update_dict(update_data=input_data)

        # ~~~ set the memory file to the flow state ~~~
        self._state_update_dict(update_data={"memory_files": self.memory_files})

        max_rounds = self.flow_config.get("max_rounds", 1)
        if max_rounds is None:
            log.info(f"Running {self.flow_config['name']} without `max_rounds` until the early exit condition is met.")

        self._sequential_run(max_rounds=max_rounds)

        output = self._get_output_from_state()

        return output