"""Struct accessor for dict/object field operations."""
from __future__ import annotations

from typing import TYPE_CHECKING

from cqla.exprs import base as exprs_base
from cqla.accessors.base import Accessor

if TYPE_CHECKING:
    from cqla.exprs.base import Expr


class Struct(Accessor):
    """Accessor for struct (dict/object) operations."""

    def _expr_repr(self) -> str:
        return f"{self._expr!r}.struct"

    def _get_dict(self, record: object) -> dict:
        val = self._expr(record)
        if isinstance(val, dict):
            return val
        return vars(val)

    def field(self, name: str, *more_names: str) -> 'Expr':
        """Extract one or more fields from the struct.

        If one name is given, returns the value.
        If multiple names are given, returns a dict with those fields.
        """
        if more_names:
            names = (name,) + more_names
            return exprs_base.Expr(
                lambda r: {n: self._get_dict(r).get(n) for n in names},
                f"{self._expr_repr()}.field({', '.join(repr(n) for n in names)})"
            )
        return exprs_base.Expr(
            lambda r: self._get_dict(r).get(name),
            f"{self._expr_repr()}.field({name!r})"
        )

    def rename_fields(self, mapping: dict[str, str]) -> 'Expr':
        """Rename fields in the struct.

        Args:
            mapping: Dict mapping old field names to new field names.

        Returns:
            A new dict with renamed fields.
        """
        def rename(r):
            d = self._get_dict(r)
            return {mapping.get(k, k): v for k, v in d.items()}

        return exprs_base.Expr(
            rename,
            f"{self._expr_repr()}.rename_fields({mapping!r})"
        )

    def with_fields(self, **exprs) -> 'Expr':
        """Add or overwrite fields in the struct.

        Args:
            **exprs: Field names mapped to values or expressions.

        Returns:
            A new dict with the added/updated fields.
        """
        def add_fields(r):
            d = dict(self._get_dict(r))
            for name, expr in exprs.items():
                if callable(expr):
                    d[name] = expr(r)
                else:
                    d[name] = expr
            return d

        exprs_repr = ', '.join(f'{k}={v!r}' for k, v in exprs.items())
        return exprs_base.Expr(
            add_fields,
            f"{self._expr_repr()}.with_fields({exprs_repr})"
        )

    def unnest(self) -> 'Expr':
        """Return the struct as a dict for merging into the parent row.

        Use in select() to expand struct fields into the row:
            .select(cq.field("metadata").struct.unnest())
        """
        return exprs_base.Expr(
            lambda r: self._get_dict(r),
            f"{self._expr_repr()}.unnest()"
        )
