from __future__ import annotations

import itertools

from typing import Callable, Iterator, Generator

import cqla.delegates as delegates
from cqla.exprs import WindowExpr, ContextExpr


type Expr = Callable[[dict | object], object]


class Query:
    def __init__(self, config: dict | object) -> None:
        self._config = config

    def filter(self, predicate: Expr) -> Query:
        # Handle ContextExpr in filter (e.g., filter(field("x").is_duplicated()))
        if isinstance(predicate, ContextExpr):
            records = list(self._config)
            lookup = predicate.compute_lookup(records)

            def gen() -> Generator:
                for i, record in enumerate(records):
                    if lookup[i]:
                        yield record
            return Query(gen())

        def gen() -> Generator:
            for record in self._config:
                if predicate(record):
                    yield record
        return Query(gen())

    def select(self, *args: Expr | str, **exprs: Expr) -> Query:

        resolved = []
        for arg in args:
            if isinstance(arg, str):
                f = delegates.field(arg)
                resolved.append((f._name, f))
            else:
                resolved.append((None, arg))

        def gen() -> Generator:
            for record in self._config:
                row = {}
                for name, expr in resolved:
                    if name is not None:
                        row[name] = expr(record)
                    else:
                        result = expr(record)
                        if isinstance(result, dict):
                            row.update(result)
                        else:
                            raise TypeError(f"Positional arg in select() must return dict, got {type(result)}")
                for name, expr in exprs.items():
                    row[name] = expr(record)
                yield row
        return Query(gen())

    def with_columns(self, **exprs: Expr) -> Query:

        window_exprs = {k: v for k, v in exprs.items() if isinstance(v, WindowExpr)}
        context_exprs = {k: v for k, v in exprs.items() if isinstance(v, ContextExpr)}
        regular_exprs = {k: v for k, v in exprs.items()
                        if not isinstance(v, (WindowExpr, ContextExpr))}

        # If any expressions need global context, buffer all records
        if window_exprs or context_exprs:
            records = list(self._config)

            window_lookups = {
                name: (expr, expr.compute_lookup(records))
                for name, expr in window_exprs.items()
            }
            context_lookups = {
                name: expr.compute_lookup(records)
                for name, expr in context_exprs.items()
            }

            def gen() -> Generator:
                for i, record in enumerate(records):
                    if isinstance(record, dict):
                        row = dict(record)
                    else:
                        row = vars(record)
                    for name, expr in regular_exprs.items():
                        row[name] = expr(record)
                    for name, (window_expr, lookup) in window_lookups.items():
                        key = window_expr._get_key(record)
                        row[name] = lookup[key]
                    for name, lookup in context_lookups.items():
                        row[name] = lookup[i]
                    yield row
            return Query(gen())
        else:
            def gen() -> Generator:
                for record in self._config:
                    if isinstance(record, dict):
                        row = dict(record)
                    else:
                        row = vars(record)
                    for name, expr in exprs.items():
                        row[name] = expr(record)
                    yield row
            return Query(gen())

    def limit(self, n: int) -> Query:
        return Query(itertools.islice(self._config, n))

    def explode(self, field_name: str) -> Query:
        """Expand a list field into multiple rows, one per element."""
        field_expr = delegates.field(field_name)

        def gen() -> Generator:
            for record in self._config:
                items = field_expr(record)
                if isinstance(record, dict):
                    for item in items:
                        row = dict(record)
                        row[field_name] = item
                        yield row
                else:
                    for item in items:
                        row = vars(record).copy()
                        row[field_name] = item
                        yield row
        return Query(gen())

    def distinct(self, *subset: str, keep: str = "first", maintain_order: bool = True) -> Query:
        """Remove duplicate rows.

        Args:
            *subset: Field names to consider for uniqueness. If empty, uses all fields.
            keep: Which duplicate to keep: 'first', 'last', 'any', 'none'.
                  'none' removes all rows that have duplicates.
            maintain_order: Whether to preserve original row order (default True).

        Returns:
            Query with duplicates removed.
        """
        def _make_hashable(value):
            if isinstance(value, dict):
                return tuple(sorted((k, _make_hashable(v)) for k, v in value.items()))
            elif isinstance(value, list):
                return tuple(_make_hashable(v) for v in value)
            elif isinstance(value, set):
                return frozenset(_make_hashable(v) for v in value)
            return value

        def _get_key(record):
            if isinstance(record, dict):
                d = record
            else:
                d = vars(record)

            if subset:
                return tuple(_make_hashable(d.get(f)) for f in subset)
            else:
                return tuple(sorted((k, _make_hashable(v)) for k, v in d.items()))

        if keep in ("first", "any"):
            def gen() -> Generator:
                seen = set()
                for record in self._config:
                    key = _get_key(record)
                    if key not in seen:
                        seen.add(key)
                        yield record
            return Query(gen())

        elif keep == "last":
            records = list(self._config)
            seen = {}
            for i, record in enumerate(records):
                key = _get_key(record)
                seen[key] = i

            if maintain_order:
                keep_indices = set(seen.values())
                def gen() -> Generator:
                    for i, record in enumerate(records):
                        if i in keep_indices:
                            yield record
                return Query(gen())
            else:
                def gen() -> Generator:
                    for idx in sorted(seen.values()):
                        yield records[idx]
                return Query(gen())

        elif keep == "none":
            records = list(self._config)
            counts: dict = {}
            for record in records:
                key = _get_key(record)
                counts[key] = counts.get(key, 0) + 1

            def gen() -> Generator:
                for record in records:
                    key = _get_key(record)
                    if counts[key] == 1:
                        yield record
            return Query(gen())

        else:
            raise ValueError(f"keep must be 'first', 'last', 'any', or 'none', got {keep!r}")

    def group_by(self, *keys: Expr | str) -> GroupByQuery:
        keys = tuple(delegates.field(k) for k in keys)

        groups = {}
        for record in self._config:
            if len(keys) == 1:
                group_key = keys[0](record)
            else:
                group_key = tuple(k(record) for k in keys)

            groups.setdefault(group_key, []).append(record)

        return GroupByQuery(groups, keys)

    def collect(self) -> list:
        return list(self._config)
    
    def __iter__(self) -> Iterator: yield from self._config

    def to_dict(self) -> dict:
        return dict(self._config)


class GroupByQuery:
    def __init__(self, groups: dict, keys: tuple[Expr, ...]) -> None:
        self._keys = keys
        self._groups = groups

    def having(self, *predicates) -> 'GroupByQuery':
        """Filter groups with predicates. Multiple predicates are combined with &."""
        filtered_groups = {}
        for group_key, records in self._groups.items():
            include = all(pred(records) for pred in predicates)
            if include:
                filtered_groups[group_key] = records
        return GroupByQuery(filtered_groups, self._keys)

    def agg(self, **exprs: Expr) -> Query:
        def gen() -> Generator:
            for group_key, records in self._groups.items():
                row = {}

                if len(self._keys) == 1:
                    key_name = getattr(self._keys[0], '_name', 'key')
                    row[key_name] = group_key
                else:
                    for key_expr, key_val in zip(self._keys, group_key):
                        key_name = getattr(key_expr, '_name', 'key')
                        row[key_name] = key_val

                # Add aggregations
                for agg_name, agg_expr in exprs.items():
                    row[agg_name] = agg_expr(records)

                yield row

        return Query(gen())
