"""Table block implementation."""

from __future__ import annotations

from typing import Any, Literal

from yera.events.blocks.base import _BlockFactory, _StreamHandle
from yera.events.models import OutputEvent, TableData
from yera.events.stream import push_output


def _is_pandas_like(data: Any) -> bool:
    """Check if data is a pandas-like object (has .columns and .values)."""
    return hasattr(data, "columns") and hasattr(data, "values")


def _is_dict_of_lists(data: Any) -> bool:
    """Check if data is a dict where all values are lists (Streamlit column format)."""
    return isinstance(data, dict) and all(
        isinstance(v, list) for v in data.values() if v is not None
    )


def _is_non_empty_list(data: Any) -> bool:
    """Check if data is a non-empty list."""
    return isinstance(data, list) and len(data) > 0


def _is_list_of_dicts(data: Any) -> bool:
    """Check if data is a non-empty list where the first element is a dict."""
    return _is_non_empty_list(data) and isinstance(data[0], dict)


def _is_list_of_lists(data: Any) -> bool:
    """Check if data is a non-empty list where the first element is also a list."""
    return _is_non_empty_list(data) and isinstance(data[0], list)


def _is_iterable_but_not_string(item: Any) -> bool:
    """Check if item is iterable but not a string."""
    return hasattr(item, "__iter__") and not isinstance(item, str)


def _generate_columns(num_cols: int) -> list[str]:
    """Generate default column names."""
    return [f"Column {i + 1}" for i in range(num_cols)]


def _pad_rows(rows: list[list]) -> tuple[list[str], list[list]]:
    """Pad rows to same length and generate column names."""
    if not rows:
        return [], []

    num_cols = max(len(row) for row in rows)
    columns = _generate_columns(num_cols)
    padded_rows = [row + [None] * (num_cols - len(row)) for row in rows]
    return columns, padded_rows


def _convert_pandas(data) -> tuple[list[str], list[list]]:
    """Convert pandas DataFrame."""
    columns = list(data.columns)
    rows = [[cell for cell in row] for row in data.values.tolist()]
    return columns, rows


def _convert_dict(data: dict) -> tuple[list[str], list[list]]:
    """Convert dict to table format."""
    columns = list(data.keys())

    if _is_dict_of_lists(data):
        # Transpose: column-oriented -> row-oriented
        max_len = max(len(v) for v in data.values() if isinstance(v, list))
        rows = [
            [data[col][i] if i < len(data[col]) else None for col in columns]
            for i in range(max_len)
        ]
        return columns, rows

    # Single row fallback
    return columns, [[data[col] for col in columns]]


def _convert_list_of_dicts(data: list) -> tuple[list[str], list[list]]:
    """Convert list of dicts."""
    columns = list(data[0].keys())
    rows = [[row.get(col) for col in columns] for row in data]
    return columns, rows


def _convert_list_of_lists(data: list) -> tuple[list[str], list[list]]:
    """Convert list of lists."""
    rows = [[cell for cell in row] for row in data]
    return _pad_rows(rows)


def _convert_iterable(data) -> tuple[list[str], list[list]]:
    """Try to convert generic iterable."""
    try:
        data_list = list(data)
        if _is_non_empty_list(data_list) and _is_iterable_but_not_string(data_list[0]):
            rows = [[cell for cell in row] for row in data_list]
            return _pad_rows(rows)
    except (TypeError, ValueError):
        pass

    raise ValueError(
        f"Unsupported data type for table: {type(data)}. "
        "Supported types: pandas DataFrame, dict with list values, "
        "list of dicts, list of lists, iterables."
    )


def _convert_data_to_table_format(
    data: Any,
) -> tuple[list[str], list[list[str | int | float]]]:
    """Convert various data formats to table format (columns, rows).

    Matches Streamlit's st.table behavior:
    - pandas DataFrame: uses .columns and .values
    - dict: keys are column names, values are lists of column data
      e.g., {"Name": ["Alice", "Bob"], "Age": [25, 30]}
    - list of dicts: keys from first dict become columns, each dict becomes a row
    - list of lists: all lists are treated as rows (no automatic header row)
    - numpy arrays, iterables: converted to rows
    """
    if data is None:
        return [], []

    if _is_pandas_like(data):
        return _convert_pandas(data)

    if isinstance(data, dict):
        return _convert_dict(data)

    if _is_list_of_dicts(data):
        return _convert_list_of_dicts(data)

    if _is_non_empty_list(data):
        if _is_list_of_lists(data):
            return _convert_list_of_lists(data)
        return [], [[cell for cell in data]]

    return _convert_iterable(data)


def _normalize_border(
    border: bool | Literal["horizontal"],
) -> Literal["all", "none", "horizontal"]:
    """Convert Streamlit-style border (bool | "horizontal") to internal string format."""
    if border is True:
        return "all"
    if border is False:
        return "none"
    # border == "horizontal"
    return "horizontal"


class TableStream(_StreamHandle):
    """Stream handle for a table block that supports incremental row additions.

    Returned by `yr.table()` to enable `add_rows()` functionality.
    Maintains table metadata (columns, border, block_id) and streams incremental
    row updates to the frontend. The frontend accumulates these rows to build the
    complete table state.

    Example:
        my_table = yr.table({"Name": ["Alice"], "Age": [25]})
        my_table.add_rows([["Bob", 30]])  # Streams only the new rows

    """

    def __init__(
        self,
        block_id: str,
        columns: list[str],
        border: bool | Literal["horizontal"],
    ):
        super().__init__(block_id, "table")
        self.columns = columns
        # Convert border to internal string format
        self.border = _normalize_border(border)

    def _emit_table_event(
        self, rows_to_send: list[list[str | int | float]]
    ) -> OutputEvent:
        """Create and add an event with the given rows."""
        event = self._create_chunk_event(
            data=TableData(columns=self.columns, rows=rows_to_send, border=self.border),
        )
        push_output(event)
        return event

    def _convert_data_to_rows_for_columns(
        self, data: Any
    ) -> list[list[str | int | float]]:
        """Convert various data formats into rows aligned to the given columns.

        Used for incremental updates (add_rows), where the column schema is
        already known from the initial table.
        """
        if _is_pandas_like(data):
            return [[cell for cell in row] for row in data.values.tolist()]

        # Handle list inputs
        if _is_non_empty_list(data):
            first = data[0]
            if _is_list_of_dicts(data):
                # List of dicts - align to existing columns
                return [[row.get(col, "") for col in self.columns] for row in data]
            if _is_list_of_lists(data):
                # List of lists - treat all as rows
                return [[cell for cell in row] for row in data]
            # Single list - treat as one row
            return [[cell for cell in data]]

        if isinstance(data, dict):
            return [[data.get(col, "") for col in self.columns]]

        raise ValueError(
            f"Unsupported data type for add_rows: {type(data)}. "
            "Supported: list of lists, list of dicts, pandas DataFrame, single list/dict."
        )

    def add_rows(self, data: Any) -> None:
        """Add rows to the table, matching Streamlit's element.add_rows() behavior.

        Supports:
        - list of lists: each inner list is a row
        - list of dicts: each dict becomes a row (keys should match existing columns)
        - pandas DataFrame: adds all rows
        - single list: treated as a single row
        """
        new_rows = self._convert_data_to_rows_for_columns(data)

        # Send only the new rows, not the entire table
        # Frontend maintains full table state by appending these rows
        self._emit_table_event(new_rows)


class _TableBlockFactory(_BlockFactory):
    """Factory for creating table blocks, matching Streamlit's st.table API."""

    def __init__(self):
        super().__init__("table")

    def __call__(
        self,
        data: Any = None,
        *,
        border: bool | Literal["horizontal"] = True,
    ) -> TableStream:
        """Create a table block, matching Streamlit's st.table() signature.

        Args:
            data: Table data. Supports (matching Streamlit):
                - pandas DataFrame: primary format
                - dict with list values: {"Name": ["Alice", "Bob"], "Age": [25, 30]}
                - list of dicts: [{"Name": "Alice", "Age": 25}, ...]
                - list of lists: all lists are rows (no header row)
                - numpy arrays, iterables: converted to rows
                - None: creates an empty table
            border: Whether to show borders. Can be True, False, or "horizontal".

        Returns:
            TableStream that supports add_rows() method, matching Streamlit's element.add_rows()

        Example:
            # Streamlit-style usage:
            my_table = yr.table({"Name": ["Alice", "Bob"], "Age": [25, 30]})
            my_table.add_rows([["Charlie", 35]])

        """
        block_id = self._generate_block_id()

        if data is None:
            columns, rows = [], []
        else:
            columns, rows = _convert_data_to_table_format(data)

        # Create table stream and immediately create the event
        stream = TableStream(block_id, columns, border)
        if columns or rows:
            stream._emit_table_event(rows)
        return stream


# Create the table callable/context manager
table = _TableBlockFactory()

__all__ = ["table"]
