from __future__ import annotations

from collections.abc import Iterable
from html import escape
from typing import Generic

from pydiverse.transform._typing import ImplT
from pydiverse.transform.core.expressions import (
    Column,
    LambdaColumn,
    SymbolicExpression,
)


class Table(Generic[ImplT]):
    """
    All attributes of a table are columns except for the `_impl` attribute
    which is a reference to the underlying table implementation.
    """

    def __init__(self, implementation: ImplT):
        self._impl = implementation

    def __getitem__(self, key) -> SymbolicExpression[Column]:
        if isinstance(key, SymbolicExpression):
            key = key._
        return SymbolicExpression(self._impl.get_col(key))

    def __setitem__(self, col, expr):
        """Mutate a column
        :param col: Either a str or SymbolicColumn
        """
        from pydiverse.transform.core.verbs import mutate

        col_name = None

        if isinstance(col, SymbolicExpression):
            underlying = col._
            if isinstance(underlying, (Column, LambdaColumn)):
                col_name = underlying.name
        elif isinstance(col, str):
            col_name = col

        if not col_name:
            raise KeyError(
                f"Invalid key {col}. Must be either a string, Column or LambdaColumn."
            )
        self._impl = (self >> mutate(**{col_name: expr}))._impl

    def __getattr__(self, name) -> SymbolicExpression[Column]:
        return SymbolicExpression(self._impl.get_col(name))

    def __iter__(self) -> Iterable[SymbolicExpression[Column]]:
        # Capture current state (this allows modifying the table inside a loop)
        cols = [
            SymbolicExpression(self._impl.get_col(name))
            for name, _ in self._impl.selected_cols()
        ]
        return iter(cols)

    def __eq__(self, other):
        return isinstance(other, type(self)) and self._impl == other._impl

    def __ne__(self, other):
        return not self.__eq__(other)

    def __hash__(self):
        return hash(self._impl)

    def __dir__(self):
        return sorted(self._impl.named_cols.fwd.keys())

    def __contains__(self, item):
        if isinstance(item, SymbolicExpression):
            item = item._
        if isinstance(item, LambdaColumn):
            return item.name in self._impl.named_cols.fwd
        if isinstance(item, Column):
            return item.uuid in self._impl.available_cols
        return False

    def __copy__(self):
        impl_copy = self._impl.copy()
        return self.__class__(impl_copy)

    def __str__(self):
        from pydiverse.transform.core.verbs import collect

        try:
            return (
                f"Table: {self._impl.name}, backend: {type(self._impl).__name__}\n"
                f"{self >> collect()}"
            )
        except Exception as e:
            return (
                f"Table: {self._impl.name}, backend: {type(self._impl).__name__}\n"
                "Failed to collect table due to an exception:\n"
                f"{type(e).__name__}: {str(e)}"
            )

    def _repr_html_(self) -> str | None:
        from pydiverse.transform.core.verbs import collect

        html = (
            f"Table <code>{self._impl.name}</code> using"
            f" <code>{type(self._impl).__name__}</code> backend:</br>"
        )
        try:
            # TODO: For lazy backend only show preview (eg. take first 20 rows)
            html += (self >> collect())._repr_html_()
        except Exception as e:
            html += (
                "</br><pre>Failed to collect table due to an exception:\n"
                f"{escape(e.__class__.__name__)}: {escape(str(e))}</pre>"
            )
        return html

    def _repr_pretty_(self, p, cycle):
        p.text(str(self) if not cycle else "...")
