from typing import Any
from pandasai.pipelines.base_logic_unit import BaseLogicUnit
from pandasai.pipelines.logic_units.code_executor import BaseCodeExecutor


class SDFCodeExecutor(BaseLogicUnit):
    """
    Executes the code generated by the prompt
    """

    def execute(self, input: Any, **kwargs) -> Any:
        code_exec = BaseCodeExecutor()
        logger = kwargs.get("logger")
        try:
            logger.log(f"Executing code: {input}")
            namespace = code_exec.execute(input, **kwargs)

            if "df" not in namespace:
                raise ValueError(f"Unable to execute code: {input}")

            return {"type": "dataframe", "value": namespace["df"]}

        except Exception:
            if logger is not None:
                logger.log("Error in executing code")
            raise
