import duckdb
from resolver import block

from blocks.base import BaseBlock

from ..patterns.data_operator import ExtendedDataFrame, ExtendedDataFrameList


@block(name="DataFrame2Str", kind="data process")
class DataFrame2Str(BaseBlock):
    """
    Convert a DataFrame to a string.


    """

    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, data: ExtendedDataFrame) -> str:
        if self._type == "markdown":
            return data.to_markdown(index=self._show_index, tablefmt="psql")
        elif self._type == "jsonstring":
            return data.to_json()
        else:
            raise NotImplementedError(
                "output_type must be either 'markdown' or 'jsonstring'"
            )


@block(name="DataFrames2StrList", kind="data process")
class DataFrames2StrList(BaseBlock):
    """
    Convert a DataFrame list to a string list.
    """

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


@block(name="DataFrame_Query_Private", kind="data process")
class DataFrameQueryPrivate(BaseBlock):
    """
    Applies a SQL-like query to a DataFrame and returns a new DataFrame with the filtered results.

    Args:
        df (ExtendedDataFrame): The input DataFrame to query.
        query (str): The query string is a duckdb SQL statement and the table name is always `result`.

    Returns:
        ExtendedDataFrame: A new DataFrame with the filtered results.
    """

    def __init__(self, query: str):
        self.query = query

    def __call__(self, df: ExtendedDataFrame) -> ExtendedDataFrame:
        result = df  # result is used as the table name in self.query
        return ExtendedDataFrame(duckdb.sql(self.query).df())