from typing import Any, Type, TypeVar

from pandasai.pipelines.base_logic_unit import BaseLogicUnit
from pandasai.responses.context import Context
from pandasai.responses.response_parser import IResponseParser, ResponseParser

IResponseParserImplementation = TypeVar(
    "IResponseParserImplementation", bound=IResponseParser
)


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

    _response_parser: Type[IResponseParserImplementation]

    def __init__(
        self, response_parser: Type[IResponseParserImplementation] = ResponseParser
    ):
        super().__init__()
        self._response_parser = response_parser

    def execute(self, input: Any, **kwargs) -> Any:
        dfs = kwargs["context"].dfs

        context = Context(kwargs["config"], kwargs["logger"], dfs[0].engine)

        return self._response_parser(context).parse(input)