"""Base expression types - no accessor imports to avoid circular dependencies."""
from __future__ import annotations

from typing import Callable, Self, TYPE_CHECKING

import cqla.delegates as delegates

if TYPE_CHECKING:
    from cqla.delegates import FieldDelegate


class AliasMixin:
    """Mixin that provides .alias() for naming expressions."""

    _name: str | None

    def alias(self, name: str) -> Self:
        clone = self.__class__.__new__(self.__class__)
        clone.__dict__.update(self.__dict__)
        clone._name = name
        return clone


class Expr(AliasMixin):
    """Base expression class - wraps a callable that extracts/computes a value from a record."""

    def __init__(self, fn: Callable[[object], object], repr_str: str | None = None) -> None:
        self._fn = fn
        self._repr = repr_str or '<lambda>'
        self._name: str | None = None

    def __repr__(self) -> str:
        return self._repr

    def __call__(self, record: object) -> object:
        return self._fn(record)

    def apply(self, fn: Callable) -> Expr:
        return Expr(lambda r: fn(self._fn(r)), f'{self!r}.apply(<fn>)')

    # Boolean operators
    def __invert__(self) -> Expr:
        return Expr(lambda r: not self._fn(r), f'~{self!r}')

    def __and__(self, other: Callable) -> Expr:
        return Expr(lambda r: self._fn(r) and other(r), f'({self!r} & {other!r})')

    def __or__(self, other: Callable) -> Expr:
        return Expr(lambda r: self._fn(r) or other(r), f'({self!r} | {other!r})')

    # Comparison operators
    def __eq__(self, other: object) -> Expr:
        return Expr(lambda r: self._fn(r) == other, f'({self._repr} == {other!r})')

    def __ne__(self, other: object) -> Expr:
        return Expr(lambda r: self._fn(r) != other, f'({self._repr} != {other!r})')

    def __gt__(self, other: object) -> Expr:
        return Expr(lambda r: self._fn(r) > other, f'({self._repr} > {other!r})')

    def __ge__(self, other: object) -> Expr:
        return Expr(lambda r: self._fn(r) >= other, f'({self._repr} >= {other!r})')

    def __lt__(self, other: object) -> Expr:
        return Expr(lambda r: self._fn(r) < other, f'({self._repr} < {other!r})')

    def __le__(self, other: object) -> Expr:
        return Expr(lambda r: self._fn(r) <= other, f'({self._repr} <= {other!r})')


class AggExpr(AliasMixin):
    """Aggregation expression - operates on a list of records, not a single record."""

    def __init__(self, expr: Callable, agg_fn: Callable[[list], object], repr_str: str) -> None:
        self._expr = expr
        self._agg_fn = agg_fn
        self._repr = repr_str
        self._name: str | None = None

    def __repr__(self) -> str:
        return self._repr

    def __call__(self, records: list) -> object:
        values = [self._expr(r) for r in records]
        return self._agg_fn(values)

    def over(self, *partition_by: str | 'FieldDelegate') -> WindowExpr:
        partition_by = tuple(delegates.field(p) for p in partition_by)
        return WindowExpr(self, partition_by)

    # Comparison operators for having() predicates
    def __eq__(self, other: object) -> 'AggPredicate':
        return AggPredicate(self, lambda a, b: a == b, other, f'({self!r} == {other!r})')

    def __ne__(self, other: object) -> 'AggPredicate':
        return AggPredicate(self, lambda a, b: a != b, other, f'({self!r} != {other!r})')

    def __gt__(self, other: object) -> 'AggPredicate':
        return AggPredicate(self, lambda a, b: a > b, other, f'({self!r} > {other!r})')

    def __ge__(self, other: object) -> 'AggPredicate':
        return AggPredicate(self, lambda a, b: a >= b, other, f'({self!r} >= {other!r})')

    def __lt__(self, other: object) -> 'AggPredicate':
        return AggPredicate(self, lambda a, b: a < b, other, f'({self!r} < {other!r})')

    def __le__(self, other: object) -> 'AggPredicate':
        return AggPredicate(self, lambda a, b: a <= b, other, f'({self!r} <= {other!r})')


class AggPredicate:
    """Predicate that compares an aggregation result - used in having()."""

    def __init__(
        self,
        agg_expr: AggExpr,
        compare_fn: Callable[[object, object], bool],
        other: object,
        repr_str: str,
    ) -> None:
        self._agg_expr = agg_expr
        self._compare_fn = compare_fn
        self._other = other
        self._repr = repr_str

    def __repr__(self) -> str:
        return self._repr

    def __and__(self, other: 'AggPredicate') -> 'AggPredicate':
        def combined(records: list) -> bool:
            return self(records) and other(records)
        return AggPredicate.__new__(AggPredicate)._init_combined(
            combined, f'({self!r} & {other!r})'
        )

    def __or__(self, other: 'AggPredicate') -> 'AggPredicate':
        def combined(records: list) -> bool:
            return self(records) or other(records)
        return AggPredicate.__new__(AggPredicate)._init_combined(
            combined, f'({self!r} | {other!r})'
        )

    def __invert__(self) -> 'AggPredicate':
        def negated(records: list) -> bool:
            return not self(records)
        return AggPredicate.__new__(AggPredicate)._init_combined(
            negated, f'~{self!r}'
        )

    def _init_combined(self, fn: Callable[[list], bool], repr_str: str) -> 'AggPredicate':
        """Initialize a combined predicate (for & | ~ operations)."""
        self._agg_expr = None  # type: ignore
        self._compare_fn = None  # type: ignore
        self._other = None
        self._repr = repr_str
        self._combined_fn = fn
        return self

    def __call__(self, records: list) -> bool:
        if hasattr(self, '_combined_fn'):
            return self._combined_fn(records)
        agg_value = self._agg_expr(records)
        return self._compare_fn(agg_value, self._other)


class WindowExpr(AliasMixin):
    """Window expression - aggregation over partitions, broadcast back to each row."""

    def __init__(self, agg_expr: AggExpr, partition_by: tuple) -> None:
        self._agg_expr = agg_expr
        self._partition_by = partition_by
        self._name: str | None = None
        partition_repr = ', '.join(repr(p) for p in partition_by)
        self._repr = f"{agg_expr!r}.over({partition_repr})"

    def __repr__(self) -> str:
        return self._repr

    def compute_lookup(self, records: list) -> dict:
        """Compute aggregation for each partition and return a lookup dict."""
        groups: dict = {}
        for record in records:
            key = self._get_key(record)
            groups.setdefault(key, []).append(record)
        return {key: self._agg_expr(group) for key, group in groups.items()}

    def _get_key(self, record: object) -> object:
        """Get the partition key for a record."""
        if len(self._partition_by) == 1:
            return self._partition_by[0](record)
        return tuple(p(record) for p in self._partition_by)


class ContextExpr(AliasMixin):
    """Context-dependent expression - needs to see all rows to compute per-row values.

    Used for operations like is_duplicated, is_unique, is_first_distinct, is_last_distinct.
    """

    def __init__(self, field_expr: 'FieldDelegate', operation: str, repr_str: str) -> None:
        self._field_expr = field_expr
        self._operation = operation
        self._repr = repr_str
        self._name: str | None = None

    def __repr__(self) -> str:
        return self._repr

    def compute_lookup(self, records: list) -> dict:
        """Compute the boolean result for each record index.

        Returns a dict mapping record index to boolean value.
        """
        # Extract all values with their indices
        values = []
        for i, record in enumerate(records):
            val = self._field_expr(record)
            values.append((i, val))

        if self._operation == "is_duplicated":
            # Count occurrences of each value
            counts: dict = {}
            for _, val in values:
                counts[val] = counts.get(val, 0) + 1
            return {i: counts[val] > 1 for i, val in values}

        elif self._operation == "is_unique":
            # Count occurrences of each value
            counts = {}
            for _, val in values:
                counts[val] = counts.get(val, 0) + 1
            return {i: counts[val] == 1 for i, val in values}

        elif self._operation == "is_first_distinct":
            # Mark first occurrence of each value
            seen: set = set()
            result = {}
            for i, val in values:
                result[i] = val not in seen
                seen.add(val)
            return result

        elif self._operation == "is_last_distinct":
            # Mark last occurrence of each value
            # First, find the last index for each value
            last_index: dict = {}
            for i, val in values:
                last_index[val] = i
            return {i: last_index[val] == i for i, val in values}

        else:
            raise ValueError(f"Unknown context operation: {self._operation}")
