from abc import ABC, abstractmethod
from typing import Any

import pandas as pd
from PIL import Image

from pandasai.exceptions import MethodNotImplementedError


class IResponseParser(ABC):
    @abstractmethod
    def parse(self, result: dict) -> Any:
        """
        Parses result from the chat input
        Args:
            result (dict): result contains type and value
        Raises:
            ValueError: if result is not a dictionary with valid key

        Returns:
            Any: Returns depending on the user input
        """
        raise MethodNotImplementedError


class ResponseParser(IResponseParser):
    _context = None

    def __init__(self, context) -> None:
        """
        Initialize the ResponseParser with Context from Agent
        Args:
            context (Context): context contains the config and logger
        """
        self._context = context

    def parse(self, result: dict) -> Any:
        """
        Parses result from the chat input
        Args:
            result (dict): result contains type and value
        Raises:
            ValueError: if result is not a dictionary with valid key

        Returns:
            Any: Returns depending on the user input
        """
        if not isinstance(result, dict) or any(
            key not in result for key in ["type", "value"]
        ):
            raise ValueError("Unsupported result format")

        if result["type"] == "plot":
            return self.format_plot(result)
        elif result["type"] == "dataframe":
            return self.format_dataframe(result)
        else:
            return result["value"]

    def format_dataframe(self, result: dict) -> Any:
        if isinstance(result["value"], dict):
            print("Df conversiont")
            df = pd.Dataframe(result["value"])
            print("Df conversiont Done")
            result["value"] = df

        return result["value"]

    def format_plot(self, result: dict) -> Any:
        """
        Display matplotlib plot against a user query.

        If `open_charts` option set to `False`, the chart won't be displayed.

        Args:
            result (dict): result contains type and value
        Returns:
            Any: Returns depending on the user input
        """
        if (
            self._context._config.open_charts
            and isinstance(result["value"], str)
            and "data:image/png;base64" not in result["value"]
        ):
            with Image.open(result["value"]) as img:
                img.show()

        return result["value"]