from typing import Any, Dict, Iterable, List, Optional

import pandas as pd
import pyarrow as pa
from duckdb import DuckDBPyRelation
from triad import Schema

from fugue import (
    ArrayDataFrame,
    ArrowDataFrame,
    DataFrame,
    LocalBoundedDataFrame,
    LocalDataFrame,
)
from fugue.exceptions import FugueDataFrameOperationError, FugueDatasetEmptyError
from fugue.plugins import (
    as_fugue_dataset,
    as_local_bounded,
    get_column_names,
    get_num_partitions,
    is_df,
)

from ._utils import encode_column_name, to_duck_type, to_pa_type


class DuckDataFrame(LocalBoundedDataFrame):
    """DataFrame that wraps DuckDB ``DuckDBPyRelation``.

    :param rel: ``DuckDBPyRelation`` object
    """

    def __init__(self, rel: DuckDBPyRelation):
        self._rel = rel
        super().__init__(schema=self._get_schema)

    def _get_schema(self) -> Schema:
        return Schema(
            [
                pa.field(x, to_pa_type(y))
                for x, y in zip(self._rel.columns, self._rel.types)
            ]
        )

    @property
    def alias(self) -> str:
        return "_" + str(id(self._rel))  # DuckDBPyRelation.alias is not always unique

    @property
    def native(self) -> DuckDBPyRelation:
        """DuckDB relation object"""
        return self._rel

    def native_as_df(self) -> DuckDBPyRelation:
        return self._rel

    @property
    def empty(self) -> bool:
        return self._rel.fetchone() is None

    def peek_array(self) -> List[Any]:
        res = self._rel.fetchone()
        if res is None:
            raise FugueDatasetEmptyError()
        return list(res)  # type: ignore

    def count(self) -> int:
        return self._rel.aggregate("count(1) AS ct").fetchone()[0]

    def _drop_cols(self, cols: List[str]) -> DataFrame:
        cols = [col for col in self._rel.columns if col not in cols]
        rel = self._rel.project(",".join(encode_column_name(n) for n in cols))
        return DuckDataFrame(rel)

    def _select_cols(self, keys: List[Any]) -> DataFrame:
        rel = self._rel.project(",".join(encode_column_name(n) for n in keys))
        return DuckDataFrame(rel)

    def rename(self, columns: Dict[str, str]) -> DataFrame:
        _assert_no_missing(self._rel, columns.keys())
        expr = ", ".join(
            f"{a} AS {b}"
            for a, b in [
                (encode_column_name(name), encode_column_name(columns.get(name, name)))
                for name in self._rel.columns
            ]
        )
        return DuckDataFrame(self._rel.project(expr))

    def alter_columns(self, columns: Any) -> DataFrame:
        new_schema = self._get_altered_schema(columns)
        if new_schema == self.schema:
            return self
        fields: List[str] = []
        for f1, f2 in zip(self.schema.fields, new_schema.fields):
            if f1.type == f2.type:
                fields.append(encode_column_name(f1.name))
            else:
                tp = to_duck_type(f2.type)
                fields.append(
                    f"CAST({encode_column_name(f1.name)} AS {tp}) "
                    f"AS {encode_column_name(f1.name)}"
                )
        return DuckDataFrame(self._rel.project(", ".join(fields)))

    def as_arrow(self, type_safe: bool = False) -> pa.Table:
        return self._rel.arrow()

    def as_pandas(self) -> pd.DataFrame:
        if any(pa.types.is_nested(f.type) for f in self.schema.fields):
            # Duckdb has issue to directly convert nested types to pandas
            return ArrowDataFrame(self.as_arrow()).as_pandas()
        return self._rel.to_df()

    def as_local(self) -> LocalDataFrame:
        return ArrowDataFrame(self.as_arrow())

    def as_array(
        self, columns: Optional[List[str]] = None, type_safe: bool = False
    ) -> List[Any]:
        if columns is not None:
            return self[columns].as_array(type_safe=type_safe)
        return self._fetchall(self._rel)

    def as_array_iterable(
        self, columns: Optional[List[str]] = None, type_safe: bool = False
    ) -> Iterable[Any]:
        if columns is not None:
            yield from self[columns].as_array_iterable(type_safe=type_safe)
        else:
            yield from self._fetchall(self._rel)

    def head(
        self, n: int, columns: Optional[List[str]] = None
    ) -> LocalBoundedDataFrame:
        if columns is not None:
            return self[columns].head(n)
        return ArrayDataFrame(self._fetchall(self._rel.limit(n)), schema=self.schema)

    def _fetchall(self, rel: DuckDBPyRelation) -> List[List[Any]]:
        map_pos = [i for i, t in enumerate(self.schema.types) if pa.types.is_map(t)]
        if len(map_pos) == 0:
            return [list(x) for x in rel.fetchall()]
        else:

            def to_list(row: Any) -> List[Any]:
                res = list(row)
                for p in map_pos:
                    res[p] = list(zip(row[p]["key"], row[p]["value"]))
                return res

            return [to_list(x) for x in rel.fetchall()]


@as_fugue_dataset.candidate(lambda df, **kwargs: isinstance(df, DuckDBPyRelation))
def _duckdb_as_fugue_df(df: DuckDBPyRelation, **kwargs: Any) -> DuckDataFrame:
    return DuckDataFrame(df, **kwargs)


@is_df.candidate(lambda df: isinstance(df, DuckDBPyRelation))
def _duck_is_df(df: DuckDBPyRelation) -> bool:
    return True


@get_num_partitions.candidate(lambda df: isinstance(df, DuckDBPyRelation))
def _duckdb_num_partitions(df: DuckDBPyRelation) -> int:
    return 1


@as_local_bounded.candidate(lambda df: isinstance(df, DuckDBPyRelation))
def _duck_as_local(df: DuckDBPyRelation) -> DuckDBPyRelation:
    return df


@get_column_names.candidate(lambda df: isinstance(df, DuckDBPyRelation))
def _get_duckdb_columns(df: DuckDBPyRelation) -> List[Any]:
    return list(df.columns)


def _assert_no_missing(df: DuckDBPyRelation, columns: Iterable[Any]) -> None:
    missing = set(columns) - set(df.columns)
    if len(missing) > 0:
        raise FugueDataFrameOperationError("found nonexistent columns: {missing}")
