ftmq.store

class PreservingLinker(Linker[StatementEntity]):
    """A linker that answers with the canonical id a statement already carries.

    Backend writers stamp `canonical_id` from the store's linker: point the store
    at this one and [`preserve`][ftmq.store.base.PreservingLinker.preserve] each
    statement before writing it to keep its canonical id. Any other id (such as
    entity-typed values) maps to itself.
    """

    def __init__(self) -> None:
        super().__init__({})
        self._stmt: Statement | None = None

    def preserve(self, stmt: Statement) -> None:
        """Resolve this statement's entity id to the canonical id it carries."""
        self._stmt = stmt

    def get_canonical(self, entity_id: str) -> str:
        if self._stmt is not None and entity_id == self._stmt.entity_id:
            return self._stmt.canonical_id
        return entity_id

class Store(nk.Store[Dataset, StatementEntity]):
    """Feature add-ons to `nomenklatura.store.Store`."""

    def __init__(
        self,
        dataset: Dataset | str | None = None,
        linker: Linker | None = None,
        cast_types: bool = True,
        **kwargs,
    ) -> None:
        """Initialize a store, use [`get_store`][ftmq.store.get_store] instead.

        Args:
            dataset: A `followthemoney.Dataset` instance to limit the scope to
            linker: A `nomenklatura.Linker` instance with linked / deduped data,
                defaults to [`get_resolver`][ftmq.store.base.get_resolver]
            cast_types: Normalize statement values on write (see
                [`ftmq.statements`][ftmq.statements])
        """
        # without a `dataset` the store spans all datasets, resolved in `scope`
        self._implicit_scope = dataset is None
        self.cast_types = cast_types
        # only the SQL family passes a `uri`; nomenklatura's stores don't take one
        uri = kwargs.pop("uri", None)
        linker = linker or get_resolver(uri)
        super().__init__(dataset=ensure_dataset(dataset), linker=linker, **kwargs)

    def writer(self, *args: Any, **kwargs: Any) -> Writer:
        """The backend writer, casting statement values on write.

        Values are cast into their property type's canonical format (see
        [`ftmq.statements`][ftmq.statements]), values that don't parse pass
        through unchanged. Disable with the store's `cast_types=False`.
        """
        return self.casting_writer(super().writer(*args, **kwargs))

    def casting_writer(self, writer: Writer) -> Writer:
        """Rebless a backend writer so it casts statement values on write.

        A store that builds its writer itself has to route it through here.
        """
        if self.cast_types:
            cls: type[Any] = type(writer)
            writer.__class__ = _casting_writer(cls)
        return writer

    def get_scope(self) -> Dataset:
        """Return the implicit `Dataset` spanning all datasets in the store."""
        raise NotImplementedError

    @property
    def scope(self) -> Dataset:
        """The read scope: the explicit `dataset`, or all datasets in the store."""
        return self.get_scope() if self._implicit_scope else self.dataset

    view_class: type["View"]

    def view(self, scope: Dataset | None = None, external: bool = False) -> "View":
        return self.view_class(self, scope or self.dataset, external=external)

    def default_view(self, external: bool = False) -> "View":
        return self.view(self.scope, external)

    def statements(self, dataset: str | Dataset | None = None) -> Statements:
        """Iterate the raw statements in this store, as they are stored.

        Unlike [`iterate`][ftmq.store.base.Store.iterate] (entity assembly
        rewrites values and synthesizes an `id` statement), a dump of these
        loads back onto its own rows. Includes external statements, unordered.
        SQL family only.

        Args:
            dataset: `Dataset` instance or name to limit scope to

        Yields:
            Generator of `followthemoney.Statement`
        """
        raise NotImplementedError

    def iterate(self, dataset: str | Dataset | None = None) -> StatementEntities:
        """Iterate all the entities, optionally limited to a dataset.

        Args:
            dataset: `Dataset` instance or name to limit scope to

        Yields:
            Generator of `nomenklatura.entity.CompositeEntity`
        """
        if dataset is not None:
            view = self.view(ensure_dataset(dataset))
        else:
            view = self.default_view()
        yield from view.entities()

class View(nk.View[Dataset, StatementEntity]):
    """Feature add-ons to `nomenklatura.store.base.View`."""

    def query(self, query: Query | None = None) -> StatementEntities:
        """Get the entities of the view, optionally filtered by a [`Query`][ftmq.Query].

        Args:
            query: The Query filter object

        Yields:
            Generator of `followthemoney.StatementEntity`
        """
        if query:
            yield from query.apply_iter(self.entities())
        else:
            yield from self.entities()

    def get_adjacents(
        self, proxies: Iterable[StatementEntity], inverted: bool | None = False
    ) -> set[StatementEntity]:
        return {
            adjacent
            for proxy in proxies
            for _, adjacent in self.get_adjacent(proxy, inverted=bool(inverted))
        }

    def stats(self, query: Query | None = None) -> DatasetStats:
        c = Collector()
        cov = c.collect_many(self.query(query))
        return cov

    def count(self, query: Query | None = None) -> int:
        return self.stats(query).entity_count or 0

    def aggregations(self, query: Query) -> AggregatorResult | None:
        if not query.aggregations:
            return
        _ = [x for x in self.query(query)]
        if query.aggregator:
            res = dict(query.aggregator.result)
            return res