/*
 * Copyright (c) "Neo4j"
 * Neo4j Sweden AB [http://neo4j.com]
 *
 * This file is part of Neo4j.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package apoc.meta;

import static apoc.util.MapUtil.map;
import static java.lang.String.format;
import static org.neo4j.token.api.TokenConstants.ANY_LABEL;
import static org.neo4j.token.api.TokenConstants.ANY_RELATIONSHIP_TYPE;

import apoc.export.util.NodesAndRelsSubGraph;
import apoc.result.GraphResult;
import apoc.result.VirtualGraph;
import apoc.result.VirtualNode;
import apoc.result.VirtualRelationship;
import apoc.util.CollectionUtils;
import apoc.util.MapUtil;
import apoc.util.Util;
import apoc.util.collection.Iterables;
import com.google.common.collect.Sets;
import java.util.AbstractMap;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Set;
import java.util.TreeMap;
import java.util.concurrent.ThreadLocalRandom;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import java.util.stream.StreamSupport;
import org.apache.commons.lang3.tuple.Pair;
import org.neo4j.cypher.export.CypherResultSubGraph;
import org.neo4j.cypher.export.DatabaseSubGraph;
import org.neo4j.cypher.export.SubGraph;
import org.neo4j.graphdb.Direction;
import org.neo4j.graphdb.Entity;
import org.neo4j.graphdb.Label;
import org.neo4j.graphdb.Node;
import org.neo4j.graphdb.Relationship;
import org.neo4j.graphdb.RelationshipType;
import org.neo4j.graphdb.ResourceIterator;
import org.neo4j.graphdb.Result;
import org.neo4j.graphdb.Transaction;
import org.neo4j.graphdb.schema.ConstraintDefinition;
import org.neo4j.graphdb.schema.ConstraintType;
import org.neo4j.graphdb.schema.IndexDefinition;
import org.neo4j.graphdb.schema.Schema;
import org.neo4j.internal.kernel.api.Read;
import org.neo4j.internal.kernel.api.TokenRead;
import org.neo4j.internal.kernel.api.procs.ProcedureCallContext;
import org.neo4j.kernel.api.KernelTransaction;
import org.neo4j.logging.Log;
import org.neo4j.procedure.Context;
import org.neo4j.procedure.Description;
import org.neo4j.procedure.Name;
import org.neo4j.procedure.NotThreadSafe;
import org.neo4j.procedure.Procedure;
import org.neo4j.procedure.UserFunction;

/**
 * The Meta class provides metadata-related operations and functions for working with Neo4j graph database.
 * It is part of the APOC (Awesome Procedures on Cypher) library.
 */
public class MetaRestricted {
    private record MetadataKey(Types type, String key) {}

    @Context
    public Transaction tx;

    @Context
    public KernelTransaction kernelTx;

    @Context
    public Transaction transaction;

    @Context
    public Log log;

    @Context
    public ProcedureCallContext procedureCallContext;

    /**
     * Represents the result of a metadata operation.
     */
    public static class MetaResult {
        @Description("The label or type name.")
        public String label;

        @Description("The property name.")
        public String property;

        @Description("The count of seen values.")
        public long count;

        @Description("If all seen values are unique.")
        public boolean unique;

        @Description("If an index exists for this property.")
        public boolean index;

        @Description("If an existence constraint exists for this property.")
        public boolean existence;

        @Description("The type represented by this row.")
        public String type;

        @Description(
                "Indicates whether the property is an array. If the type column is \"RELATIONSHIP,\" this will be true if there is at least one node with two outgoing relationships of the type specified by the label or property column.")
        public boolean array;

        @Description("This is always null.")
        public List<Object> sample;

        @Description(
                "The ratio (rounded down) of the count of outgoing relationships for a specific label and relationship type relative to the total count of those patterns.")
        public long left; // 0,1,

        @Description(
                "The ratio (rounded down) of the count of incoming relationships for a specific label and relationship type relative to the total count of those patterns.")
        public long right; // 0,1,many

        @Description("The labels of connect nodes.")
        public List<String> other = new ArrayList<>();

        @Description(
                "For uniqueness constraints, this field shows other labels present on nodes that also contain the uniqueness constraint.")
        public List<String> otherLabels = new ArrayList<>();

        @Description("Whether this refers to a node or a relationship.")
        public String elementType;
    }

    /**
     * Represents a specific metadata item, extending MetaResult.
     */
    public static class MetaItem extends MetaResult {
        public long leftCount; // 0,1,
        public long rightCount; // 0,1,many

        public MetaItem addLabel(String label) {
            this.otherLabels.add(label);
            return this;
        }

        public MetaItem(String label, String name) {
            this.label = label;
            this.property = name;
        }

        public MetaItem inc() {
            count++;
            return this;
        }

        public MetaItem rel(long out, long in) {
            this.type = Types.RELATIONSHIP.name();
            if (out > 1) array = true;
            leftCount += out;
            rightCount += in;
            left = leftCount / count;
            right = rightCount / count;
            return this;
        }

        public MetaItem other(List<String> labels) {
            for (String l : labels) {
                if (!this.other.contains(l)) this.other.add(l);
            }
            return this;
        }

        public MetaItem type(String type) {
            this.type = type;
            return this;
        }

        public MetaItem array(boolean array) {
            this.array = array;
            return this;
        }

        public MetaItem elementType(String elementType) {
            switch (elementType) {
                case "NODE":
                    this.elementType = "node";
                    break;
                case "RELATIONSHIP":
                    this.elementType = "relationship";
                    break;
            }
            return this;
        }
    }

    /**
     * The MetaStats class represents metadata statistics collected from the transactional database.
     * It includes counts for labels, relationship types, property keys, nodes, relationships, and various maps for stats.
     */
    public static class MetaStats {
        @Description("The total number of distinct node labels.")
        public final long labelCount;

        @Description("The total number of distinct relationship types.")
        public final long relTypeCount;

        @Description("The count of property keys.")
        public final long propertyKeyCount;

        @Description("The total number of nodes.")
        public final long nodeCount;

        @Description("The total number of relationships.")
        public final long relCount;

        @Description("A map of labels to their count.")
        public final Map<String, Long> labels;

        @Description("A map of relationship types per start or end node label.")
        public final Map<String, Long> relTypes;

        @Description("A map of relationship types to their count.")
        public final Map<String, Long> relTypesCount;

        @Description("A map containing all the given return fields from this procedure.")
        public final Map<String, Object> stats;

        /**
         * Constructs a MetaStats object with the provided metadata statistics.
         *
         * @param labelCount       The count of labels in the database.
         * @param relTypeCount     The count of relationship types in the database.
         * @param propertyKeyCount The count of property keys in the database.
         * @param nodeCount        The count of nodes in the database.
         * @param relCount         The count of relationships in the database.
         * @param labels           A map of label names and their corresponding counts.
         * @param relTypes         A map of relationship type names and their corresponding counts.
         * @param relTypesCount    A map of relationship type names and their total count.
         */
        public MetaStats(
                long labelCount,
                long relTypeCount,
                long propertyKeyCount,
                long nodeCount,
                long relCount,
                Map<String, Long> labels,
                Map<String, Long> relTypes,
                Map<String, Long> relTypesCount) {
            this.labelCount = labelCount;
            this.relTypeCount = relTypeCount;
            this.propertyKeyCount = propertyKeyCount;
            this.nodeCount = nodeCount;
            this.relCount = relCount;
            this.labels = labels;
            this.relTypes = relTypes;
            this.relTypesCount = relTypesCount;
            this.stats = map(
                    "labelCount",
                    labelCount,
                    "relTypeCount",
                    relTypeCount,
                    "propertyKeyCount",
                    propertyKeyCount,
                    "nodeCount",
                    nodeCount,
                    "relCount",
                    relCount,
                    "labels",
                    labels,
                    "relTypes",
                    relTypes);
        }
    }

    /**
     * The StatsCallback interface defines callback methods for collecting label and relationship statistics.
     */
    interface StatsCallback {
        void label(String labelName, long count);

        void rel(String typeName, long count);

        void rel(String typeName, String labelName, long out, long in);
    }

    @NotThreadSafe
    @Procedure("apoc.meta.stats")
    @Description("Returns the metadata stored in the transactional database statistics.")
    public Stream<MetaStats> stats() {
        return Stream.of(collectStats());
    }

    @NotThreadSafe
    @UserFunction(name = "apoc.meta.nodes.count")
    @Description("Returns the sum of the `NODE` values with the given labels in the `LIST<STRING>`.")
    public long count(
            @Name(value = "nodes", defaultValue = "[]", description = "A list of node labels.") List<String> nodes,
            @Name(
                            value = "config",
                            defaultValue = "{}",
                            description =
                                    "A relationship, node or map to get the property types from. { includeRels = [] :: LIST<STRING> }")
                    Map<String, Object> config) {
        MetaConfig conf = new MetaConfig(config);
        final var subGraph = DatabaseSubGraph.optimizedForCount(transaction, kernelTx);
        Stream<Label> labels = CollectionUtils.isEmpty(nodes)
                ? StreamSupport.stream(subGraph.getAllLabelsInUse().spliterator(), false)
                : nodes.stream().filter(Objects::nonNull).map(String::trim).map(Label::label);

        final boolean isIncludeRels = CollectionUtils.isEmpty(conf.getIncludeRels());
        Set<String> visitedNodes = new HashSet<>();
        return labels.flatMap(label -> isIncludeRels
                        ? Stream.of(subGraph.countsForNode(label))
                        : conf.getIncludeRels().stream()
                                .filter(Objects::nonNull)
                                .map(String::trim)
                                .map(rel -> {
                                    final int lastCharIdx = rel.length() - 1;
                                    final Direction direction;
                                    switch (rel.charAt(lastCharIdx)) {
                                        case '>':
                                            direction = Direction.OUTGOING;
                                            rel = rel.substring(0, lastCharIdx);
                                            break;
                                        case '<':
                                            direction = Direction.INCOMING;
                                            rel = rel.substring(0, lastCharIdx);
                                            break;
                                        default:
                                            direction = Direction.BOTH;
                                    }
                                    return Pair.of(direction, rel);
                                })
                                .flatMap(pair -> transaction
                                        .findNodes(label)
                                        .map(node -> {
                                            if (!visitedNodes.contains(node.getElementId())
                                                    && node.hasRelationship(
                                                            pair.getLeft(),
                                                            RelationshipType.withName(pair.getRight()))) {
                                                visitedNodes.add(node.getElementId());
                                                return 1L;
                                            } else {
                                                return 0L;
                                            }
                                        })
                                        .stream()))
                .reduce(0L, Math::addExact);
    }

    /**
     * Collects statistics about the graph database.
     *
     * @return The collected statistics about the graph database.
     */
    private MetaStats collectStats() {
        Map<String, Long> relStatsCount = new LinkedHashMap<>();
        TokenRead tokenRead = kernelTx.tokenRead();
        Read read = kernelTx.dataRead();

        long relTypeCount = Iterables.count(tx.getAllRelationshipTypesInUse());
        long labelCount = Iterables.count(tx.getAllLabelsInUse());

        Map<String, Long> labelStats = new LinkedHashMap<>((int) labelCount);
        Map<String, Long> relStats = new LinkedHashMap<>(2 * (int) relTypeCount);

        collectStats(DatabaseSubGraph.optimizedForCount(transaction, kernelTx), new StatsCallback() {
            @Override
            public void label(String labelName, long count) {
                if (count > 0) labelStats.put(labelName, count);
            }

            @Override
            public void rel(String typeName, long count) {
                if (count > 0) {
                    relStatsCount.merge(typeName, count, Long::sum);
                    relStats.put("()-[:" + typeName + "]->()", count);
                }
            }

            @Override
            public void rel(String typeName, String labelName, long out, long in) {
                if (out > 0) {
                    relStats.put("(:" + labelName + ")-[:" + typeName + "]->()", out);
                }
                if (in > 0) {
                    relStats.put("()-[:" + typeName + "]->(:" + labelName + ")", in);
                }
            }
        });

        return new MetaStats(
                labelCount,
                relTypeCount,
                tokenRead.propertyKeyCount(),
                read.countsForNode(ANY_LABEL),
                read.countsForRelationship(ANY_LABEL, ANY_RELATIONSHIP_TYPE, ANY_LABEL),
                labelStats,
                relStats,
                relStatsCount);
    }

    /**
     * Collects statistics about a subgraph of the database.
     *
     * @param subGraph      The subgraph to collect statistics for.
     * @param cb            The callback to receive the collected statistics.
     */
    private void collectStats(SubGraph subGraph, StatsCallback cb) {

        final var labels = subGraph.getAllLabelsInUse();
        final var types = subGraph.getAllRelationshipTypesInUse();

        labels.forEach(label -> {
            long count = subGraph.countsForNode(label);
            if (count > 0) {
                String name = label.name();
                cb.label(name, count);
                types.forEach(type -> {
                    long relCountOut = subGraph.countsForRelationship(label, type);
                    long relCountIn = subGraph.countsForRelationship(type, label);
                    cb.rel(type.name(), name, relCountOut, relCountIn);
                });
            }
        });
        types.forEach(type -> {
            String name = type.name();
            cb.rel(name, subGraph.countsForRelationship(type));
        });
    }

    // todo ask index for distinct values if index size < 10 or so
    // todo put index sizes for indexed properties
    @NotThreadSafe
    @Procedure("apoc.meta.data")
    @Description("Examines the full graph and returns a table of metadata.")
    public Stream<MetaResult> data(
            @Name(
                            value = "config",
                            defaultValue = "{}",
                            description =
                                    "Number of nodes to sample, setting sample to `-1` will remove sampling; { sample = 1000 :: INTEGER }")
                    Map<String, Object> config) {
        SampleMetaConfig metaConfig = new SampleMetaConfig(config);
        return collectMetaData(DatabaseSubGraph.optimizedForCount(transaction, kernelTx), metaConfig).values().stream()
                .flatMap(x -> x.values().stream());
    }

    public record MetaMapResult(@Description("Meta information represented as a map.") Map<String, Object> value) {}

    @NotThreadSafe
    @Procedure("apoc.meta.schema")
    @Description("Examines the given sub-graph and returns metadata as a `MAP`.")
    public Stream<MetaMapResult> schema(
            @Name(
                            value = "config",
                            defaultValue = "{}",
                            description =
                                    "Number of nodes to sample, setting sample to `-1` will remove sampling; { sample = 1000 :: INTEGER }")
                    Map<String, Object> config) {
        MetaStats metaStats = collectStats();
        SampleMetaConfig metaConfig = new SampleMetaConfig(config);
        Map<MetadataKey, Map<String, MetaItem>> metaData =
                collectMetaData(DatabaseSubGraph.optimizedForCount(transaction, kernelTx), metaConfig);

        Map<String, Object> relationships = collectRelationshipsMetaData(metaStats, metaData);
        Map<String, Object> nodes = collectNodesMetaData(metaStats, metaData, relationships);
        final Collection<String> commonKeys = Sets.intersection(nodes.keySet(), relationships.keySet());
        if (!commonKeys.isEmpty()) {
            relationships = relationships.entrySet().stream()
                    .map(e -> {
                        final String key = e.getKey();
                        return commonKeys.contains(key)
                                ? new AbstractMap.SimpleEntry<>(
                                        format("%s (%s)", key, Types.RELATIONSHIP.name()), e.getValue())
                                : e;
                    })
                    .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
        }
        nodes.putAll(relationships);
        return Stream.of(new MetaMapResult(nodes));
    }

    /**
     * This procedure is intended to replicate what's in the core Neo4j product, but with the crucial difference that it
     * supports flexible sampling options, and does not scan the entire database.  The result is producing a table of
     * metadata that is useful for generating "Tables 4 Labels" schema designs for RDBMSs, but in a more performant way.
     */
    @NotThreadSafe
    @Procedure("apoc.meta.nodeTypeProperties")
    @Description(
            "Examines the full graph and returns a table of metadata with information about the `NODE` values therein.")
    public Stream<Tables4LabelsProfile.NodeTypePropertiesEntry> nodeTypeProperties(
            @Name(
                            value = "config",
                            defaultValue = "{}",
                            description =
                                    """
            {
                    includeLabels = [] :: LIST<STRING>,
                    includeRels = [] :: LIST<STRING>,
                    excludeLabels = [] :: LIST<STRING>,
                    excludeRels = [] :: LIST<STRING>,
                    sample = 1000 :: INTEGER,
                    maxRels = 100 :: INTEGER
            }
            """)
                    Map<String, Object> config) {
        MetaConfig metaConfig = new MetaConfig(config);
        try {
            return collectTables4LabelsProfile(metaConfig).asNodeStream();
        } catch (Exception e) {
            log.debug("apoc.meta.nodeTypeProperties(): Failed to return stream", e);
            throw new RuntimeException(e);
        }
    }

    /**
     * This procedure is intended to replicate what's in the core Neo4j product, but with the crucial difference that it supports flexible sampling options, and
     * does not scan the entire database.  The result is producing a table of metadata that is useful for generating "Tables 4 Labels" schema designs for
     * RDBMSs, but in a more performant way.
     */
    @NotThreadSafe
    @Procedure("apoc.meta.relTypeProperties")
    @Description(
            "Examines the full graph and returns a table of metadata with information about the `RELATIONSHIP` values therein.")
    public Stream<Tables4LabelsProfile.RelTypePropertiesEntry> relTypeProperties(
            @Name(
                            value = "config",
                            defaultValue = "{}",
                            description =
                                    """
            {
                    includeLabels = [] :: LIST<STRING>,
                    includeRels = [] :: LIST<STRING>,
                    excludeLabels = [] :: LIST<STRING>,
                    excludeRels = [] :: LIST<STRING>,
                    sample = 1000 :: INTEGER,
                    maxRels = 100 :: INTEGER
            }
            """)
                    Map<String, Object> config) {
        MetaConfig metaConfig = new MetaConfig(config);
        try {
            return collectTables4LabelsProfile(metaConfig).asRelStream();
        } catch (Exception e) {
            log.debug("apoc.meta.relTypeProperties(): Failed to return stream", e);
            throw new RuntimeException(e);
        }
    }

    /**
     * Collects metadata for generating a "Tables 4 Labels" profile based on the provided configuration. This method
     * examines the graph schema and counts the occurrences of labels and relationships. It also observes nodes and their
     * relationships, applying sampling and filtering based on the configuration. The collected metadata is stored in a
     * Tables4LabelsProfile object.
     */
    private Tables4LabelsProfile collectTables4LabelsProfile(MetaConfig config) {
        Tables4LabelsProfile profile = new Tables4LabelsProfile();

        // Collect constraints from the schema
        Schema schema = tx.schema();

        for (ConstraintDefinition cd : schema.getConstraints()) {
            if (cd.isConstraintType(ConstraintType.NODE_PROPERTY_EXISTENCE)) {
                List<String> props = new ArrayList<>(10);
                if (ConstraintTracker.nodeConstraints.containsKey(cd.getLabel().name())) {
                    props = ConstraintTracker.nodeConstraints.get(cd.getLabel().name());
                }
                cd.getPropertyKeys().forEach(props::add);
                ConstraintTracker.nodeConstraints.put(cd.getLabel().name(), props);

            } else if (cd.isConstraintType(ConstraintType.RELATIONSHIP_PROPERTY_EXISTENCE)) {
                List<String> props = new ArrayList<>(10);
                if (ConstraintTracker.relConstraints.containsKey(
                        cd.getRelationshipType().name())) {
                    props = ConstraintTracker.relConstraints.get(
                            cd.getRelationshipType().name());
                }
                cd.getPropertyKeys().forEach(props::add);
                ConstraintTracker.relConstraints.put(cd.getRelationshipType().name(), props);
            }
        }

        // Get label counts from the label count store
        Map<String, Long> countStore = getLabelCountStore(transaction, kernelTx);

        Set<String> includeLabels = config.getIncludeLabels();
        Set<String> excludeLabels = config.getExcludeLabels();

        Set<String> includeRels = config.getIncludeRels();
        Set<String> excludeRels = config.getExcludeRels();

        // Iterate over labels in the graph
        for (Label label : tx.getAllLabelsInUse()) {
            String labelName = label.name();

            if (!excludeLabels.contains(labelName) && (includeLabels.isEmpty() || includeLabels.contains(labelName))) {
                // Skip if explicitly excluded or at least 1 include specified and not included

                long labelCount = countStore.get(labelName);
                long sample = getSampleForLabelCount(labelCount, config.getSample());

                try (ResourceIterator<Node> nodes = tx.findNodes(label)) {
                    int count = 1;
                    while (nodes.hasNext()) {
                        Node node = nodes.next();
                        if (count++ % sample == 0) {
                            final Set<Boolean> skips = StreamSupport
                                    // we analyze the node for each its relationship type
                                    .stream(node.getRelationshipTypes().spliterator(), false)
                                    .map(rel -> excludeRels.contains(
                                                    rel.name()) // we skip a node when the user said that must be
                                            // excluded
                                            // or when the user provided and inclusion list, but it's not in the
                                            // provided list
                                            || (!includeRels.isEmpty() && !includeRels.contains(rel.name())))
                                    .collect(Collectors.toSet());
                            // if the Set has just one element and is true we skip the node
                            // if there are two elements [true, false] we don't skip it as give it means that
                            // it have a relationship that satisfies the condition provided
                            // by the configuration
                            if (skips.size() == 1 && skips.iterator().next()) continue;
                            profile.observe(node, config);
                        }
                    }
                }
            }
        }

        return profile.finished();
    }

    // End new code
    /**
     * Collects metadata for generating a metadata map based on the provided subgraph and configuration. This method iterates
     * over the labels and relationships in the subgraph, collects various metadata information, and stores it in the
     * metadata map.
     */
    private Map<MetadataKey, Map<String, MetaItem>> collectMetaData(SubGraph graph, SampleMetaConfig config) {
        Map<MetadataKey, Map<String, MetaItem>> metaData = new LinkedHashMap<>(100);

        Set<RelationshipType> types = Iterables.asSet(graph.getAllRelationshipTypesInUse());
        Map<String, Iterable<ConstraintDefinition>> relConstraints = new HashMap<>(20);
        Map<String, Set<String>> relIndexes = new HashMap<>();
        for (RelationshipType type : graph.getAllRelationshipTypesInUse()) {
            metaData.put(new MetadataKey(Types.RELATIONSHIP, type.name()), new LinkedHashMap<>(10));
            relConstraints.put(type.name(), graph.getConstraints(type));
            relIndexes.put(type.name(), getIndexedProperties(graph.getIndexes(type)));
        }
        for (Label label : graph.getAllLabelsInUse()) {
            Map<String, MetaItem> nodeMeta = new LinkedHashMap<>(50);
            String labelName = label.name();
            // workaround in case of duplicated keys
            metaData.put(new MetadataKey(Types.NODE, labelName), nodeMeta);
            Iterable<ConstraintDefinition> constraints = graph.getConstraints(label);
            Set<String> indexed = getIndexedProperties(graph.getIndexes(label));
            long labelCount = graph.countsForNode(label);
            long sample = getSampleForLabelCount(labelCount, config.getSample());
            Iterator<Node> nodes = graph.findNodes(label);
            int count = 1;
            while (nodes.hasNext()) {
                Node node = nodes.next();
                if (count++ % sample == 0) {
                    addRelationships(metaData, nodeMeta, labelName, node, relConstraints, types, relIndexes);
                    addProperties(nodeMeta, labelName, constraints, indexed, node, node);
                }
            }
        }
        return metaData;
    }

    private Set<String> getIndexedProperties(Iterable<IndexDefinition> indexes) {
        return Iterables.stream(indexes)
                .map(IndexDefinition::getPropertyKeys)
                .flatMap(Iterables::stream)
                .collect(Collectors.toSet());
    }

    private static Map<String, Long> getLabelCountStore(Transaction tx, KernelTransaction kernelTx) {
        List<String> labels =
                Iterables.stream(tx.getAllLabelsInUse()).map(Label::name).collect(Collectors.toList());
        TokenRead tokenRead = kernelTx.tokenRead();
        return labels.stream()
                .collect(Collectors.toMap(e -> e, e -> kernelTx.dataRead().countsForNode(tokenRead.nodeLabel(e))));
    }

    public static long getSampleForLabelCount(long labelCount, long sample) {
        if (sample != -1L) {
            long skipCount = labelCount / sample;
            long min = (long) Math.floor(skipCount - (skipCount * 0.1D));
            long max = (long) Math.ceil(skipCount + (skipCount * 0.1D));
            if (min >= max) {
                return -1L;
            }
            long randomValue = ThreadLocalRandom.current().nextLong(min, max);
            return randomValue == 0L ? -1L : randomValue; // it can't return zero as it's used in % ops
        } else {
            return sample;
        }
    }

    private Map<String, Object> collectNodesMetaData(
            MetaStats metaStats, Map<MetadataKey, Map<String, MetaItem>> metaData, Map<String, Object> relationships) {
        Map<String, Object> nodes = new LinkedHashMap<>();
        Map<String, List<Map<String, Object>>> startNodeNameToRelationshipsMap = new HashMap<>();
        for (MetadataKey metadataKey : metaData.keySet()) {
            Map<String, MetaItem> entityData = metaData.get(metadataKey);
            Map<String, Object> entityProperties = new LinkedHashMap<>();
            Map<String, Object> entityRelationships = new LinkedHashMap<>();
            List<String> labels = new LinkedList<>();
            boolean isNode = metaStats.labels.keySet().stream().anyMatch(metadataKey.key::equals);
            for (String entityDataKey : entityData.keySet()) {
                MetaItem metaItem = entityData.get(entityDataKey);
                if (metaItem.elementType.equals("relationship")) {
                    isNode = false;
                    break;
                } else {
                    if (metaItem.unique) labels = metaItem.otherLabels;
                    if (!metaItem.type.equals("RELATIONSHIP")) { // NODE PROPERTY
                        entityProperties.put(
                                entityDataKey,
                                MapUtil.map(
                                        "type",
