from __future__ import annotations

from dataclasses import dataclass
from datetime import date, datetime, time
from decimal import Decimal
from math import floor
from typing import Any, Literal, cast

import narwhals as nw

SortDirection = Literal["asc", "desc"]
FilterOperator = Literal[
    "contains",
    "starts_with",
    "ends_with",
    "eq",
    "neq",
    "gt",
    "gte",
    "lt",
    "lte",
    "between",
]
Frame = nw.DataFrame[Any] | nw.LazyFrame[Any]


@dataclass(frozen=True)
class SortSpec:
    column: str
    direction: SortDirection = "asc"

    @property
    def descending(self) -> bool:
        return self.direction == "desc"


@dataclass(frozen=True)
class FilterSpec:
    column: str
    operator: FilterOperator
    value: Any
    value_to: Any | None = None


class TableEngine:
    """Narwhals-backed query engine for the tablescope widget."""

    def __init__(
        self,
        source: Any,
        *,
        page_size: int = 100,
        frequency_limit: int = 25,
    ) -> None:
        normalized = nw.from_native(source, eager_only=False)
        if not isinstance(normalized, (nw.DataFrame, nw.LazyFrame)):
            msg = (
                "tablescope expected an object that narwhals can treat as a "
                f"dataframe or lazy frame, got {type(source)!r}."
            )
            raise TypeError(msg)

        self._frame: Frame = normalized
        self._page_size = max(1, page_size)
        self._frequency_limit = max(1, frequency_limit)
        self._total_rows_cache: int | None = None

        schema = (
            normalized.collect_schema()
            if isinstance(normalized, nw.LazyFrame)
            else normalized.schema
        )
        self._schema = dict(schema.items())
        self._columns = list(self._schema)
        self._supports_random_access = self._can_paginate_frame(normalized)

    @property
    def page_size(self) -> int:
        return self._page_size

    @property
    def supports_random_access(self) -> bool:
        return self._supports_random_access

    @property
    def columns(self) -> list[str]:
        return list(self._columns)

    @property
    def total_rows(self) -> int:
        if self._total_rows_cache is None:
            result = self._materialize(
                self._as_lazy(self._frame).select(nw.len().alias("__total_rows__"))
            )
            row = self._rows_from_frame(result)[0]
            self._total_rows_cache = int(row["__total_rows__"])
        return self._total_rows_cache

    def initial_payload(self) -> dict[str, Any]:
        return {
            "columns": [self._column_payload(name) for name in self._columns],
            "page_size": self._page_size,
            "total_rows": self.total_rows,
            "supports_random_access": self._supports_random_access,
            "rows": self.fetch_rows(
                offset=0,
                limit=self._page_size,
                sort=[],
                filters=[],
            ),
        }

    def fetch_rows(
        self,
        *,
        offset: int,
        limit: int | None = None,
        sort: list[SortSpec] | None = None,
        filters: list[FilterSpec] | None = None,
    ) -> dict[str, Any]:
        filter_specs = filters or []
        filtered = self._apply_filters(self._frame, filter_specs)
        total_rows = self._frame_total_rows(filtered)
        bounded_offset = max(0, min(offset, max(0, total_rows - 1)))
        effective_limit = max(1, limit or self._page_size)
        bounded_limit = min(effective_limit, max(1, total_rows - bounded_offset))
        sort_specs = sort or []

        transformed = self._apply_sort(filtered, sort_specs)
        page = self._page_frame(transformed, offset=bounded_offset, limit=bounded_limit)
        rows = self._rows_from_frame(self._materialize(page))

        for index, row in enumerate(rows):
            row["__tablescope_row_number__"] = bounded_offset + index

        return {
            "offset": bounded_offset,
            "limit": bounded_limit,
            "total_rows": total_rows,
            "sort": [
                {"column": spec.column, "direction": spec.direction}
                for spec in sort_specs
            ],
            "filters": [
                {
                    "column": spec.column,
                    "operator": spec.operator,
                    "value": spec.value,
                    "value_to": spec.value_to,
                }
                for spec in filter_specs
            ],
            "rows": rows,
        }

    def jump_to_fraction(
        self,
        fraction: float,
        *,
        limit: int | None = None,
        sort: list[SortSpec] | None = None,
        filters: list[FilterSpec] | None = None,
    ) -> dict[str, Any]:
        filter_specs = filters or []
        total_rows = self._frame_total_rows(self._apply_filters(self._frame, filter_specs))
        effective_limit = limit or self._page_size
        if total_rows <= effective_limit:
            offset = 0
        else:
            max_offset = total_rows - effective_limit
            clamped_fraction = min(1.0, max(0.0, fraction))
            offset = floor(max_offset * clamped_fraction)
        return self.fetch_rows(
            offset=offset,
            limit=effective_limit,
            sort=sort,
            filters=filter_specs,
        )

    def describe_column(
        self,
        column: str,
        *,
        filters: list[FilterSpec] | None = None,
    ) -> dict[str, Any]:
        self._require_column(column)
        dtype = self._schema[column]
        filtered = self._apply_filters(self._frame, filters or [])
        exprs: list[Any] = [
            nw.len().alias("row_count"),
            nw.col(column).count().alias("non_null_count"),
            nw.col(column).null_count().alias("null_count"),
            nw.col(column).n_unique().alias("distinct_count"),
        ]

        if not dtype.is_nested():
            exprs.extend(
                [
                    nw.col(column).min().alias("min"),
                    nw.col(column).max().alias("max"),
                ]
            )

        if dtype.is_numeric():
            exprs.extend(
                [
                    nw.col(column).mean().alias("mean"),
                    nw.col(column).median().alias("median"),
                    nw.col(column).std(ddof=1).alias("std"),
                    nw.col(column).sum().alias("sum"),
                ]
            )

        frame = self._materialize(self._as_lazy(filtered).select(*exprs))
        values = self._rows_from_frame(frame)[0]
        return {
            "column": column,
            "dtype": str(dtype),
            "is_numeric": dtype.is_numeric(),
            "is_temporal": dtype.is_temporal(),
            "stats": values,
        }

    def value_counts(
        self,
        column: str,
        *,
        limit: int | None = None,
        filters: list[FilterSpec] | None = None,
    ) -> dict[str, Any]:
        self._require_column(column)
        dtype = self._schema[column]
        if dtype.is_nested():
            raise TypeError(f"Column {column!r} is nested and cannot be grouped.")

        top_n = max(1, limit or self._frequency_limit)
        lazy = (
            self._as_lazy(self._apply_filters(self._frame, filters or []))
            .group_by(column)
            .agg(nw.len().alias("count"))
            .sort(["count", column], descending=[True, False])
            .head(top_n)
        )
        frame = self._materialize(lazy)
        return {
            "column": column,
            "limit": top_n,
            "values": self._rows_from_frame(frame),
        }

    def _column_payload(self, name: str) -> dict[str, Any]:
        dtype = self._schema[name]
        return {
            "name": name,
            "dtype": str(dtype),
            "is_numeric": dtype.is_numeric(),
            "is_temporal": dtype.is_temporal(),
            "is_nested": dtype.is_nested(),
        }

    def _require_column(self, column: str) -> None:
        if column not in self._schema:
            raise KeyError(f"Unknown column {column!r}.")

    def _apply_filters(self, frame: Frame, filters: list[FilterSpec]) -> Frame:
        result = frame
        for spec in filters:
            self._require_column(spec.column)
            result = result.filter(self._filter_expr(spec))
        return result

    def _apply_sort(self, frame: Frame, sort_specs: list[SortSpec]) -> Frame:
        if not sort_specs:
            return frame

        for spec in sort_specs:
            self._require_column(spec.column)

        by = [spec.column for spec in sort_specs]
        descending = [spec.descending for spec in sort_specs]
        if len(by) == 1:
            return frame.sort(by[0], descending=descending[0], nulls_last=True)
        return frame.sort(by, descending=descending, nulls_last=True)

    def _frame_total_rows(self, frame: Frame) -> int:
        if frame is self._frame:
            return self.total_rows

        result = self._materialize(
            self._as_lazy(frame).select(nw.len().alias("__total_rows__"))
        )
        row = self._rows_from_frame(result)[0]
        return int(row["__total_rows__"])

    def _filter_expr(self, spec: FilterSpec) -> Any:
        dtype = self._schema[spec.column]
        expr = nw.col(spec.column)

        if spec.operator in {"contains", "starts_with", "ends_with"}:
            value = str(spec.value)
            if spec.operator == "contains":
                return expr.str.contains(value, literal=True)
            if spec.operator == "starts_with":
                return expr.str.starts_with(value)
            return expr.str.ends_with(value)

        value = self._coerce_filter_value(spec.column, spec.value)
        literal = nw.lit(value)

        if spec.operator == "eq":
            return expr == literal
        if spec.operator == "neq":
            return expr != literal
        if spec.operator == "gt":
            return expr > literal
        if spec.operator == "gte":
            return expr >= literal
        if spec.operator == "lt":
            return expr < literal
        if spec.operator == "lte":
            return expr <= literal
        if spec.operator == "between":
            if spec.value_to is None:
                raise ValueError(
                    f"Filter {spec.column!r} with operator 'between' requires value_to."
                )
            upper = nw.lit(self._coerce_filter_value(spec.column, spec.value_to))
            return (expr >= literal) & (expr <= upper)

        raise ValueError(f"Unsupported filter operator {spec.operator!r} for {dtype}.")

    def _coerce_filter_value(self, column: str, value: Any) -> Any:
        dtype = self._schema[column]
        if dtype.is_numeric():
            text = str(value).strip()
            if text == "":
                raise ValueError(f"Filter for {column!r} requires a numeric value.")
            if dtype.is_integer() or dtype.is_signed_integer() or dtype.is_unsigned_integer():
                return int(text)
            return float(text)

        if dtype.is_temporal():
            text = str(value).strip()
            if text == "":
                raise ValueError(f"Filter for {column!r} requires a temporal value.")
            dtype_name = str(dtype).lower()
            if dtype_name == "date":
                return date.fromisoformat(text)
            # datetime-local inputs may omit seconds (e.g. "2023-10-10T13:01").
            if "T" in text and text.count(":") == 1:
                text += ":00"
            return datetime.fromisoformat(text)

        return value

    def _page_frame(self, frame: Frame, *, offset: int, limit: int) -> Frame:
        if isinstance(frame, nw.DataFrame):
            arrow = frame.to_arrow()
            return cast(Frame, nw.from_native(arrow.slice(offset, limit), eager_only=True))

        native = frame.to_native()
        paged_native = self._slice_native(native, offset=offset, limit=limit)
        paged_frame = nw.from_native(paged_native, eager_only=False)
        if not isinstance(paged_frame, (nw.DataFrame, nw.LazyFrame)):
            msg = f"Unexpected paged frame type {type(paged_frame)!r}."
            raise TypeError(msg)
        return paged_frame

    def _slice_native(self, native: Any, *, offset: int, limit: int) -> Any:
        if hasattr(native, "slice"):
            try:
                return native.slice(offset, limit)
            except TypeError:
                try:
                    return native.slice(offset=offset, length=limit)
                except TypeError:
                    pass

        if hasattr(native, "limit"):
            try:
                return native.limit(limit, offset=offset)
            except TypeError:
                try:
                    return native.limit(limit, offset)
                except TypeError:
                    pass

        if offset == 0:
            return self._as_lazy(self._frame).head(limit).to_native()

        msg = (
            "This backend does not expose a native paginated slice API. "
            "tablescope can render the first page, but random-access jumps are "
            "not supported for this source yet."
        )
        raise NotImplementedError(msg)

    def _can_paginate_frame(self, frame: Frame) -> bool:
        if isinstance(frame, nw.DataFrame):
            return True
        native = frame.to_native()
        return hasattr(native, "slice") or hasattr(native, "limit")

    def _as_lazy(self, frame: Frame) -> nw.LazyFrame[Any]:
        if isinstance(frame, nw.LazyFrame):
            return frame
        return frame.lazy()

    def _materialize(self, frame: Frame) -> nw.DataFrame[Any]:
        if isinstance(frame, nw.DataFrame):
            return frame
        return frame.collect()

    def _rows_from_frame(self, frame: nw.DataFrame[Any]) -> list[dict[str, Any]]:
        rows = frame.to_arrow().to_pylist()
        return [self._jsonify_mapping(row) for row in rows]

    def _jsonify_mapping(self, value: dict[str, Any]) -> dict[str, Any]:
        return {key: self._jsonify_scalar(item) for key, item in value.items()}

    def _jsonify_scalar(self, value: Any) -> Any:
        if isinstance(value, (datetime, date, time)):
            return value.isoformat()
        if isinstance(value, Decimal):
            return str(value)
        if isinstance(value, bytes):
            return value.hex()
        if isinstance(value, dict):
            return {key: self._jsonify_scalar(item) for key, item in value.items()}
        if isinstance(value, list):
            return [self._jsonify_scalar(item) for item in value]
        return value
