from __future__ import annotations

from abc import ABC, abstractmethod
from contextlib import suppress
from datetime import date, datetime
from pathlib import Path
from typing import (
    Any,
    Dict,
    Generic,
    Iterable,
    Literal,
    Mapping,
    Sequence,
    TypeVar,
)

import pyarrow as pa
import pyarrow.compute as pc
import pyarrow.lib as pal
import pyarrow.parquet as pq
import pyarrow.types as pt
from rich.console import Console

from textual_fastdatatable.format import measure_width

AutoBackendType = Any

try:
    import polars as pl
    import polars.datatypes as pld
except ImportError:
    _HAS_POLARS = False
else:
    _HAS_POLARS = True


def create_backend(
    data: "AutoBackendType",
    max_rows: int | None = None,
    has_header: bool = False,
) -> DataTableBackend:
    if isinstance(data, pa.Table):
        return ArrowBackend(data, max_rows=max_rows)
    if isinstance(data, pa.RecordBatch):
        return ArrowBackend.from_batches(data, max_rows=max_rows)
    if _HAS_POLARS and isinstance(data, pl.DataFrame):
        return PolarsBackend.from_dataframe(data, max_rows=max_rows)

    if isinstance(data, Path) or isinstance(data, str):
        data = Path(data)
        if data.suffix in [".pqt", ".parquet"]:
            return ArrowBackend.from_parquet(data, max_rows=max_rows)
        if _HAS_POLARS:
            return PolarsBackend.from_file_path(
                data, max_rows=max_rows, has_header=has_header
            )
    if isinstance(data, Sequence) and not data:
        return ArrowBackend(pa.table([]), max_rows=max_rows)
    if isinstance(data, Sequence) and _is_iterable(data[0]):
        return ArrowBackend.from_records(data, max_rows=max_rows, has_header=has_header)

    if (
        isinstance(data, Mapping)
        and isinstance(next(iter(data.keys())), str)
        and isinstance(next(iter(data.values())), Sequence)
    ):
        return ArrowBackend.from_pydict(data, max_rows=max_rows)

    raise TypeError(
        f"Cannot automatically create backend for data of type: {type(data)}. "
        f"Data must be of type: Union[pa.Table, pa.RecordBatch, Path, str, "
        "Sequence[Iterable[Any]], Mapping[str, Sequence[Any]], pl.DataFrame",
    )


def _is_iterable(item: Any) -> bool:
    try:
        iter(item)
    except TypeError:
        return False
    else:
        return True


_TableTypeT = TypeVar("_TableTypeT")


class DataTableBackend(ABC, Generic[_TableTypeT]):
    data: _TableTypeT

    @abstractmethod
    def __init__(self, data: _TableTypeT, max_rows: int | None = None) -> None:
        pass

    @classmethod
    @abstractmethod
    def from_pydict(
        cls, data: Mapping[str, Sequence[Any]], max_rows: int | None = None
    ) -> "DataTableBackend":
        pass

    @property
    @abstractmethod
    def source_data(self) -> _TableTypeT:
        """
        Return the source data as an Arrow table
        """
        pass

    @property
    @abstractmethod
    def source_row_count(self) -> int:
        """
        The number of rows in the source data, before filtering down to max_rows
        """
        pass

    @property
    @abstractmethod
    def row_count(self) -> int:
        """
        The number of rows in backend's retained data, after filtering down to max_rows
        """
        pass

    @property
    def column_count(self) -> int:
        return len(self.columns)

    @property
    @abstractmethod
    def columns(self) -> Sequence[str]:
        """
        A list of column labels
        """
        pass

    @property
    @abstractmethod
    def column_content_widths(self) -> Sequence[int]:
        """
        A list of integers corresponding to the widest utf8 string length
        of any data in each column.
        """
        pass

    @abstractmethod
    def get_row_at(self, index: int) -> Sequence[Any]:
        pass

    @abstractmethod
    def get_column_at(self, index: int) -> Sequence[Any]:
        pass

    @abstractmethod
    def get_cell_at(self, row_index: int, column_index: int) -> Any:
        pass

    @abstractmethod
    def append_column(self, label: str, default: Any | None = None) -> int:
        """
        Returns column index
        """

    @abstractmethod
    def append_rows(self, records: Iterable[Iterable[Any]]) -> list[int]:
        """
        Returns new row indicies
        """
        pass

    @abstractmethod
    def drop_row(self, row_index: int) -> None:
        pass

    @abstractmethod
    def update_cell(self, row_index: int, column_index: int, value: Any) -> None:
        """
        Raises IndexError if bad indicies
        """

    @abstractmethod
    def sort(
        self, by: list[tuple[str, Literal["ascending", "descending"]]] | str
    ) -> None:
        """
        by: str sorts table by the data in the column with that name (asc).
        by: list[tuple] sorts the table by the named column(s) with the directions
            indicated.
        """


class ArrowBackend(DataTableBackend[pa.Table]):
    def __init__(self, data: pa.Table, max_rows: int | None = None) -> None:
        self._source_data = data

        # Arrow allows duplicate field names, but a table's to_pylist() and
        # to_pydict() methods will drop duplicate-named fields!
        field_names: list[str] = []
        renamed = False
        for field in data.column_names:
            n = 0
            while field in field_names:
                field = f"{field}{n}"
                renamed = True
                n += 1
            field_names.append(field)
        if renamed:
            data = data.rename_columns(field_names)

        self._source_row_count = data.num_rows
        if max_rows is not None and max_rows < self._source_row_count:
            self.data = data.slice(offset=0, length=max_rows)
        else:
            self.data = data
        self._console = Console()
        self._column_content_widths: list[int] = []

    @staticmethod
    def _pydict_from_records(
        records: Sequence[Iterable[Any]], has_header: bool = False
    ) -> dict[str, list[Any]]:
        headers = (
            records[0]
            if has_header
            else [f"f{i}" for i in range(len(list(records[0])))]
        )
        data = list(map(list, records[1:] if has_header else records))
        pydict = {header: [row[i] for row in data] for i, header in enumerate(headers)}
        return pydict

    @staticmethod
    def _handle_overflow(scalar: pa.Scalar) -> Any | None:
        """
        PyArrow may throw an OverflowError when casting arrow types
        to python types; in some cases we can catch these and
        present a sensible value in the data table; otherwise
        we return None.
        """
        if pt.is_date32(scalar.type):
            if scalar.value > 0:  # type: ignore[attr-defined]
                return date.max
            elif scalar.value <= 0:  # type: ignore[attr-defined]
                return date.min
        elif pt.is_date64(scalar.type):
            if scalar.value > 0:  # type: ignore[attr-defined]
                return date.max
            elif scalar.value <= 0:  # type: ignore[attr-defined]
                return date.min
        elif pt.is_timestamp(scalar.type):
            if scalar.value > 0:  # type: ignore[attr-defined]
                return datetime.max
            elif scalar.value <= 0:  # type: ignore[attr-defined]
                return datetime.min

        return None

    @classmethod
    def from_batches(
        cls, data: pa.RecordBatch, max_rows: int | None = None
    ) -> "ArrowBackend":
        tbl = pa.Table.from_batches([data])
        return cls(tbl, max_rows=max_rows)

    @classmethod
    def from_parquet(
        cls, path: Path | str, max_rows: int | None = None
    ) -> "ArrowBackend":
        tbl = pq.read_table(str(path))
        return cls(tbl, max_rows=max_rows)

    @classmethod
    def from_pydict(
        cls, data: Mapping[str, Sequence[Any]], max_rows: int | None = None
    ) -> "ArrowBackend":
        try:
            tbl = pa.Table.from_pydict(dict(data))
        except (pal.ArrowInvalid, pal.ArrowTypeError):
            # one or more fields has mixed types, like int and
            # string. Cast all to string for safety
            new_data = {
                k: [str(val) if val is not None else None for val in v]
                for k, v in data.items()
            }
            tbl = pa.Table.from_pydict(new_data)
        return cls(tbl, max_rows=max_rows)

    @classmethod
    def from_records(
        cls,
        records: Sequence[Iterable[Any]],
        has_header: bool = False,
        max_rows: int | None = None,
    ) -> "ArrowBackend":
        pydict = cls._pydict_from_records(records, has_header)
        return cls.from_pydict(pydict, max_rows=max_rows)

    @property
    def source_data(self) -> pa.Table:
        return self._source_data

    @property
    def source_row_count(self) -> int:
        return self._source_row_count

    @property
    def row_count(self) -> int:
        return self.data.num_rows

    @property
    def column_count(self) -> int:
        return self.data.num_columns

    @property
    def columns(self) -> Sequence[str]:
        return self.data.column_names

    @property
    def column_content_widths(self) -> list[int]:
        if not self._column_content_widths:
            measurements = [self._measure(arr) for arr in self.data.columns]
            # pc.max returns None for each column without rows; we need to return 0
            # instead.
            self._column_content_widths = [cw or 0 for cw in measurements]

        return self._column_content_widths

    def get_row_at(self, index: int) -> Sequence[Any]:
        try:
            row: Dict[str, Any] = self.data.slice(index, length=1).to_pylist()[0]
        except OverflowError:
            return [
                self._handle_overflow(self.data[i][index])
                for i in range(len(self.columns))
            ]
        else:
            return list(row.values())

    def get_column_at(self, column_index: int) -> list[Any]:
        try:
            values = self.data[column_index].to_pylist()
        except OverflowError:
            # TODO: consider registering a scalar UDF here for parallel processing
            return [self._handle_overflow(scalar) for scalar in self.data[column_index]]
        else:
            return values

    def get_cell_at(self, row_index: int, column_index: int) -> Any:
        scalar = self.data[column_index][row_index]
        try:
            value = scalar.as_py()
        except OverflowError:
            value = self._handle_overflow(scalar)
        return value

    def append_column(self, label: str, default: Any | None = None) -> int:
        """
        Returns column index
        """
        if default is None:
            arr: pa.Array = pa.nulls(self.row_count)
        else:
            arr = pa.nulls(self.row_count, type=pa.string())
            arr = arr.fill_null(str(default))

        self.data = self.data.append_column(label, arr)
        if self._column_content_widths:
            self._column_content_widths.append(measure_width(default, self._console))
        return self.data.num_columns - 1

    def append_rows(self, records: Iterable[Iterable[Any]]) -> list[int]:
        rows = list(records)
        indicies = list(range(self.row_count, self.row_count + len(rows)))
        records_with_headers = [self.data.column_names, *rows]
        pydict = self._pydict_from_records(records_with_headers, has_header=True)
        old_rows = self.data.to_batches()
        new_rows = pa.RecordBatch.from_pydict(
            pydict,
            schema=self.data.schema,
        )
        self.data = pa.Table.from_batches([*old_rows, new_rows])
        self._reset_content_widths()
        return indicies

    def drop_row(self, row_index: int) -> None:
        if row_index < 0 or row_index >= self.row_count:
            raise IndexError(f"Can't drop row {row_index} of {self.row_count}")
        above = self.data.slice(0, row_index).to_batches()
        below = self.data.slice(row_index + 1).to_batches()
        self.data = pa.Table.from_batches([*above, *below])
        self._reset_content_widths()
        pass

    def update_cell(self, row_index: int, column_index: int, value: Any) -> None:
        column = self.data.column(column_index)
        pycolumn = self.get_column_at(column_index=column_index)
        pycolumn[row_index] = value
        new_type = pa.string() if pt.is_null(column.type) else column.type
        self.data = self.data.set_column(
            column_index,
            self.data.column_names[column_index],
            pa.array(pycolumn, type=new_type),
        )
        if self._column_content_widths:
            self._column_content_widths[column_index] = max(
                measure_width(value, self._console),
                self._column_content_widths[column_index],
            )

    def sort(
        self, by: list[tuple[str, Literal["ascending", "descending"]]] | str
    ) -> None:
        """
        by: str sorts table by the data in the column with that name (asc).
        by: list[tuple] sorts the table by the named column(s) with the directions
            indicated.
        """
        self.data = self.data.sort_by(by)

    def _reset_content_widths(self) -> None:
        self._column_content_widths = []

    def _measure(self, arr: pa._PandasConvertible) -> int:
        # with some types we can measure the width more efficiently
        if pt.is_boolean(arr.type):
            return 7
        elif pt.is_null(arr.type):
            return 0
        elif (
            pt.is_integer(arr.type)
            or pt.is_floating(arr.type)
            or pt.is_decimal(arr.type)
        ):
            try:
                col_max = pc.max(arr.fill_null(0)).as_py()
            except OverflowError:
                col_max = 9223372036854775807
            try:
                col_min = pc.min(arr.fill_null(0)).as_py()
            except OverflowError:
                col_min = -9223372036854775807
            return max([measure_width(el, self._console) for el in [col_max, col_min]])
        elif pt.is_temporal(arr.type):
            try:
                value = arr.drop_null()[0].as_py()
            except OverflowError:
                return 26  # need space for the infinity sign and a space
            except IndexError:
                return 24
            else:
                # valid temporal types all have the same width for their type
                return measure_width(value, self._console)

        # for everything else, we need to compute it
        # First, cast the data to strings
        try:
            arr = arr.cast(
                pa.string(),
                safe=False,
            )
        except (pal.ArrowNotImplementedError, pal.ArrowInvalid):
            # some types can't be casted to strings natively by arrow, but they
            # can be casted to strings by python. The arrow way is faster, but
            # if it fails, register a python udf and try again
            def py_str(_ctx: Any, arr: pa.Array) -> str | pa.Array | pa.ChunkedArray:
                return pa.array([str(el) for el in arr], type=pa.string())

            udf_name = f"tfdt_pystr_{arr.type}"
            with suppress(pal.ArrowKeyError):  # already registered
                pc.register_scalar_function(
                    py_str,
                    function_name=udf_name,
                    function_doc={"summary": "str", "description": "built-in str"},
                    in_types={"arr": arr.type},
                    out_type=pa.string(),
                )

            arr = pc.call_function(udf_name, [arr])

        # next, try to measure the UTF-encoded string length of each cell,
        # then take the max
        try:
            width: int = pc.max(pc.utf8_length(arr.fill_null("")).fill_null(0)).as_py()
        except OverflowError:
            width = 10
        return width


if _HAS_POLARS:

    class PolarsBackend(DataTableBackend[pl.DataFrame]):
        @classmethod
        def from_file_path(
            cls, path: Path, max_rows: int | None = None, has_header: bool = True
        ) -> "PolarsBackend":
            if path.suffix in [".arrow", ".feather"]:
                tbl = pl.read_ipc(path)
            elif path.suffix == ".arrows":
                tbl = pl.read_ipc_stream(path)
            elif path.suffix == ".json":
                tbl = pl.read_json(path)
            elif path.suffix == ".csv":
                tbl = pl.read_csv(path, has_header=has_header)
            else:
                raise TypeError(
                    f"Dont know how to load file type {path.suffix} for {path}"
                )
            return cls(tbl, max_rows=max_rows)

        @classmethod
        def from_pydict(
            cls, pydict: Mapping[str, Sequence[Any]], max_rows: int | None = None
        ) -> "PolarsBackend":
            return cls(pl.from_dict(pydict), max_rows=max_rows)

        @classmethod
        def from_dataframe(
            cls, frame: pl.DataFrame, max_rows: int | None = None
        ) -> "PolarsBackend":
            return cls(frame, max_rows=max_rows)

        def __init__(self, data: pl.DataFrame, max_rows: int | None = None) -> None:
            self._source_data = data

            # Arrow allows duplicate field names, but a table's to_pylist() and
            # to_pydict() methods will drop duplicate-named fields!
            field_names: list[str] = []
            for field in data.columns:
                n = 0
                while field in field_names:
                    field = f"{field}{n}"
                    n += 1
                field_names.append(field)
            data.columns = field_names

            self._source_row_count = len(data)
            if max_rows is not None and max_rows < self._source_row_count:
                self.data = data.slice(offset=0, length=max_rows)
            else:
                self.data = data
            self._console = Console()
            self._column_content_widths: list[int] = []

        @property
        def source_data(self) -> pl.DataFrame:
            return self._source_data

        @property
        def source_row_count(self) -> int:
            return self._source_row_count

        @property
        def row_count(self) -> int:
            return len(self.data)

        @property
        def column_count(self) -> int:
            return len(self.data.columns)

        @property
        def columns(self) -> Sequence[str]:
            return self.data.columns

        def get_row_at(self, index: int) -> Sequence[Any]:
            if index < 0 or index >= len(self.data):
                raise IndexError(
                    f"Cannot get row={index} in table with {len(self.data)} rows "
                    f"and {len(self.data.columns)} cols"
                )
            return list(self.data.slice(index, length=1).to_dicts()[0].values())

        def get_column_at(self, column_index: int) -> Sequence[Any]:
            if column_index < 0 or column_index >= len(self.data.columns):
                raise IndexError(
                    f"Cannot get column={column_index} in table with {len(self.data)} "
                    f"rows and {len(self.data.columns)} cols."
                )
            return list(self.data.to_series(column_index))

        def get_cell_at(self, row_index: int, column_index: int) -> Any:
            if (
                row_index >= len(self.data)
                or row_index < 0
                or column_index < 0
                or column_index >= len(self.data.columns)
            ):
                raise IndexError(
                    f"Cannot get cell at row={row_index} col={column_index} in table "
                    f"with {len(self.data)} rows and {len(self.data.columns)} cols"
                )
            return self.data.to_series(column_index)[row_index]

        def drop_row(self, row_index: int) -> None:
            if row_index < 0 or row_index >= self.row_count:
                raise IndexError(f"Can't drop row {row_index} of {self.row_count}")
            above = self.data.slice(0, row_index)
            below = self.data.slice(row_index + 1)
            self.data = pl.concat([above, below])
            self._reset_content_widths()

        def append_rows(self, records: Iterable[Iterable[Any]]) -> list[int]:
            rows_to_add = pl.from_dicts(
                [dict(zip(self.data.columns, row)) for row in records]
            )
            indicies = list(range(self.row_count, self.row_count + len(rows_to_add)))
            self.data = pl.concat([self.data, rows_to_add])
            self._reset_content_widths()
            return indicies

        def append_column(self, label: str, default: Any | None = None) -> int:
            """
            Returns column index
            """
            self.data = self.data.with_columns(
                pl.Series([default])
                .extend_constant(default, self.row_count - 1)
                .alias(label)
            )
            if self._column_content_widths:
                self._column_content_widths.append(
                    measure_width(default, self._console)
                )
            return len(self.data.columns) - 1

        def _reset_content_widths(self) -> None:
            self._column_content_widths = []

        def update_cell(self, row_index: int, column_index: int, value: Any) -> None:
            if row_index >= len(self.data) or column_index >= len(self.data.columns):
                raise IndexError(
                    f"Cannot update cell at row={row_index} col={column_index} in "
                    f"table with {len(self.data)} rows and "
                    f"{len(self.data.columns)} cols"
                )
            col_name = self.data.columns[column_index]
            self.data = self.data.with_columns(
                self.data.to_series(column_index)
                .scatter(row_index, value)
                .alias(col_name)
            )
            if self._column_content_widths:
                self._column_content_widths[column_index] = max(
                    measure_width(value, self._console),
                    self._column_content_widths[column_index],
                )

        @property
        def column_content_widths(self) -> list[int]:
            if not self._column_content_widths:
                measurements = [
                    self._measure(self.data[arr]) for arr in self.data.columns
                ]
                # pc.max returns None for each column without rows; we need to return 0
                # instead.
                self._column_content_widths = [cw or 0 for cw in measurements]

            return self._column_content_widths

        def _measure(self, arr: pl.Series) -> int:
            # with some types we can measure the width more efficiently
            dtype = arr.dtype
            if dtype == pld.Categorical():
                return self._measure(arr.cat.get_categories())

            if dtype.is_decimal() or dtype.is_float() or dtype.is_integer():
                col_max = arr.max()
                col_min = arr.min()
                return max(
                    [measure_width(el, self._console) for el in [col_max, col_min]]
                )
            if dtype.is_temporal():
                try:
                    value = arr.drop_nulls()[0]
                except IndexError:
                    return 0
                else:
                    return measure_width(value, self._console)
            if dtype.is_(pld.Boolean()):
                return 7

            # for everything else, we need to compute it

            arr = arr.cast(
                pl.Utf8(),
                strict=False,
            )
            width = arr.fill_null("<null>").str.len_chars().max()
            assert isinstance(width, int)
            return width

        def sort(
            self, by: list[tuple[str, Literal["ascending", "descending"]]] | str
        ) -> None:
            """
            by: str sorts table by the data in the column with that name (asc).
            by: list[tuple] sorts the table by the named column(s) with the directions
                indicated.
            """
            if isinstance(by, str):
                cols = [by]
                typs = [False]
            else:
                cols = [x for x, _ in by]
                typs = [x == "descending" for _, x in by]
            self.data = self.data.sort(cols, descending=typs)
