import json
from abc import ABC, abstractmethod

from langchain.chat_models.base import BaseChatModel
from langchain.prompts import StringPromptTemplate
from pandas import DataFrame
from pandasai import SmartDataframe
from pandasai.llm import LangchainLLM

from utils.usage import with_openai_stats


class ExtendedDataFrame(DataFrame):
    def __init__(
        self, data=None, index=None, columns=None, dtype=None, copy=False
    ):  # df type in pandas bool | None is Unresitered
        super().__init__(data, index, columns, dtype, copy)

    def __str__(self) -> str:
        return super().__str__()

    def json_serialize(self) -> str:
        # return json.loads(self.to_json())
        return self.to_markdown()


class ExtendedDataFrameList(ABC):
    def __init__(self, list_of_df: list) -> None:
        assert all(isinstance(df, ExtendedDataFrame) for df in list_of_df)
        self.list_of_df = list_of_df

    def __str__(self) -> str:
        return json.dumps([df.json_serialize() for df in self.list_of_df])

    def json_serialize(self) -> list:
        return [df.json_serialize() for df in self.list_of_df]


class DataFrameOperator(ABC):
    def __init__(self):
        ...

    @abstractmethod
    def __call__(self, input: ExtendedDataFrame) -> ExtendedDataFrame:
        ...


class DescribeOperator(DataFrameOperator):
    def __call__(self, input: ExtendedDataFrame) -> ExtendedDataFrame:
        return ExtendedDataFrame(input.describe())


class IsMonotonicIncreaseOperator(DataFrameOperator):
    def __call__(self, input: ExtendedDataFrame) -> ExtendedDataFrame:
        input.interpolate(method="linear", inplace=True)
        is_column_monotonic_inc = input.apply(lambda x: x.is_monotonic_increasing)
        return ExtendedDataFrame(is_column_monotonic_inc)


class AIDataFrameOperator(DataFrameOperator):
    def __init__(self, model: BaseChatModel, prompt_template: StringPromptTemplate):
        # the custom_whitelisted_dependencies param is not supportted yet.
        custom_whitelisted_dependencies = [""]

        # convert langchain llm to openai llm because pandas-ai only support langchain 0.0.199
        model = LangchainLLM(model)

        self.llm = model
        self.custom_whitelisted_dependencies = custom_whitelisted_dependencies
        self.prompt_template = prompt_template

    @with_openai_stats
    def __call__(self, input: ExtendedDataFrame, **kwargs) -> ExtendedDataFrame:
        sdf = SmartDataframe(
            input,
            config={
                "llm": self.llm,
                "verbose": False,
                "enable_cache": False,
                "custom_whitelisted_dependencies": self.custom_whitelisted_dependencies,
            },
        )
        prompt = self.prompt_template.format(**kwargs)

        response = sdf.chat(prompt, output_type="dataframe")
        if isinstance(response, SmartDataframe):
            return ExtendedDataFrame(response.dataframe)
        else:
            raise TypeError(f"response is not a DataFrame but {response}")