import json

from classes.data_operator import ExtendedDataFrame, ExtendedDataFrameList
from classes.derived_str import JsonString
from scheduler.block import BaseBlock


class List2ConversationText(BaseBlock):
    def __call__(self, messages: list) -> str:
        ms = []
        for i, m in enumerate(messages):
            role = "asistant"
            if i % 2 == 0:
                role = "user"
            ms.append(f"{role}: {m}")
        return "\n".join(ms)


class JoinList(BaseBlock):
    def __init__(self, template: str, delimiter: str = "\n"):
        super(JoinList, self).__init__()
        self.template = template
        self.delimiter = delimiter

    def __call__(self, **kwargs) -> str:
        if len(kwargs) == 0:
            return None

        xs = []
        for t in zip(*kwargs.values()):
            xs.append(self.template.format(**dict(zip(kwargs.keys(), t))))
        return self.delimiter.join(xs)


class KeySelector(BaseBlock):
    def __init__(self, key: str):
        super(KeySelector).__init__()
        self.key = key

    def __call__(self, messages: dict) -> str:
        return messages.get(self.key)


class DataFrame2Str(BaseBlock):
    def __init__(self, output_type: str = "markdown", show_index: bool = False) -> None:
        if not output_type in ["markdown", "jsonstring"]:
            raise TypeError(
                f"output_type must be either 'markdown' or 'jsonstring', but got {output_type}"
            )
        super(DataFrame2Str, self).__init__()
        self._type = output_type
        self._show_index = show_index

    def __call__(self, input: ExtendedDataFrame) -> str:
        if self._type == "markdown":
            return input.to_markdown(index=self._show_index, tablefmt="psql")
        elif self._type == "jsonstring":
            return input.to_json()
        else:
            raise NotImplementedError(
                "output_type must be either 'markdown' or 'jsonstring'"
            )


class DataFrames2StrList(BaseBlock):
    def __init__(self, output_type: str = "markdown", show_index: bool = False) -> None:
        if not output_type in ["markdown", "jsonstring"]:
            raise TypeError(
                f"output_type must be either 'markdown' or 'jsonstring', but got {output_type}"
            )
        super(DataFrames2StrList, self).__init__()
        self._type = output_type
        self._show_index = show_index

    def __call__(self, input: ExtendedDataFrameList) -> list:
        if self._type == "markdown":
            return [
                df.to_markdown(index=self._show_index, tablefmt="psql")
                for df in input.list_of_df
            ]
        elif self._type == "jsonstring":
            return [df.to_json() for df in input.list_of_df]
        else:
            raise NotImplementedError(
                f"output_type must be either 'markdown' or 'jsonstring'"
            )


class Dict2JsonString(BaseBlock):
    def __call__(self, input_obj: dict) -> JsonString:
        return JsonString(json.dumps(input_obj))


class KeySelector4JsonString(BaseBlock):
    def __init__(self, key: str):
        super(KeySelector4JsonString).__init__()
        self.key = key

    def __call__(self, input: JsonString) -> JsonString:
        messages = input.parsed()
        try:
            content = messages.get(self.key)
            if isinstance(content, str):
                return JsonString(content)
            return JsonString(json.dumps(content))
        except Exception as e:
            raise AttributeError(
                f"The input could not be parsed as on object. Input: {input}. Error message: {e}"
            )