"""
Copyright 2024, Zep Software, Inc.

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.
"""

import logging
from typing import Any

from graphiti_core.driver.driver import GraphProvider
from graphiti_core.driver.falkordb import STOPWORDS
from graphiti_core.driver.operations.search_ops import SearchOperations
from graphiti_core.driver.query_executor import QueryExecutor
from graphiti_core.driver.record_parsers import (
    community_node_from_record,
    entity_edge_from_record,
    entity_node_from_record,
    episodic_node_from_record,
)
from graphiti_core.edges import EntityEdge
from graphiti_core.graph_queries import (
    get_nodes_query,
    get_relationships_query,
    get_vector_cosine_func_query,
)
from graphiti_core.models.edges.edge_db_queries import get_entity_edge_return_query
from graphiti_core.models.nodes.node_db_queries import (
    COMMUNITY_NODE_RETURN,
    EPISODIC_NODE_RETURN,
    get_entity_node_return_query,
)
from graphiti_core.nodes import CommunityNode, EntityNode, EpisodicNode
from graphiti_core.search.search_filters import (
    SearchFilters,
    edge_search_filter_query_constructor,
    node_search_filter_query_constructor,
)

logger = logging.getLogger(__name__)

MAX_QUERY_LENGTH = 128

# ---------------------------------------------------------------------------
# Vector index dispatcher (PATCHED 2026-05-02, BirdAI vendored patch).
#
# graphiti-core's FalkorDB driver historically composed similarity queries
# using `vec.cosineDistance(...)` in interpreted Cypher, which produces a
# full-table scan for every search. FalkorDB supports native vector indexes
# via `db.idx.vector.queryNodes` and `db.idx.vector.queryRelationships`;
# this dispatcher uses them when present and falls back to the cosine math
# otherwise.
#
# Index existence is checked once per (label, attribute, entity_type) and
# cached at module scope. The cache should be invalidated whenever
# `build_indices_and_constraints` runs (since indexes may have been created
# or dropped). FalkorDriver.build_indices_and_constraints is patched to
# call `_invalidate_falkordb_vector_index_cache()` after building.
#
# Over-fetch factor (VECTOR_INDEX_CANDIDATE_MULTIPLIER from graph_queries)
# preserves recall when WHERE filters reject some of the top-k candidates.
# ---------------------------------------------------------------------------

from graphiti_core.graph_queries import (
    VECTOR_INDEX_CANDIDATE_MULTIPLIER,
    get_vector_cosine_func_query,
)

# Cache: key = (label, attribute, entity_type), value = bool
# entity_type is 'NODE' or 'RELATIONSHIP'.
_FALKORDB_VECTOR_INDEX_CACHE: dict[tuple[str, str, str], bool] = {}


def _invalidate_falkordb_vector_index_cache() -> None:
    """Clear the vector-index existence cache. Call after build_indices_and_constraints."""
    _FALKORDB_VECTOR_INDEX_CACHE.clear()


async def _falkordb_vector_index_exists(
    executor: QueryExecutor,
    label: str,
    attribute: str,
    entity_type: str,
) -> bool:
    """Check whether a FalkorDB vector index exists for the given target.

    entity_type is 'NODE' for node-label indexes, 'RELATIONSHIP' for edge-type indexes.
    Result is cached at module scope; call _invalidate_falkordb_vector_index_cache()
    after building or dropping indexes.
    """
    key = (label, attribute, entity_type)
    if key in _FALKORDB_VECTOR_INDEX_CACHE:
        return _FALKORDB_VECTOR_INDEX_CACHE[key]

    try:
        records, _, _ = await executor.execute_query(
            "CALL db.indexes() YIELD label, properties, types, entitytype "
            "RETURN label, properties, types, entitytype"
        )
    except Exception as e:
        # If we cannot enumerate indexes, fall back to "no index" rather than
        # propagating the error. The fallback cosine-math path is correct,
        # just slower.
        logger.warning(f"FalkorDB vector index probe failed; assuming none exist: {e}")
        _FALKORDB_VECTOR_INDEX_CACHE[key] = False
        return False

    found = False
    for r in records:
        # Records come back as dict-like rows keyed by column name (not
        # tuples). Access by string keys matching the YIELD clause above.
        rec_label = r.get('label') if hasattr(r, 'get') else r['label']
        rec_props = r.get('properties') if hasattr(r, 'get') else r['properties']
        rec_types = r.get('types') if hasattr(r, 'get') else r['types']
        rec_entitytype = r.get('entitytype') if hasattr(r, 'get') else r['entitytype']
        if rec_props is None:
            rec_props = []
        if rec_types is None:
            rec_types = {}

        if rec_label != label:
            continue
        if rec_entitytype is not None and rec_entitytype != entity_type:
            continue
        if attribute not in rec_props:
            continue

        # rec_types is a dict like {attribute: ['VECTOR', ...], ...} or sometimes
        # a flat list — handle both shapes.
        if isinstance(rec_types, dict):
            attr_types = rec_types.get(attribute, [])
        else:
            attr_types = rec_types
        if 'VECTOR' in attr_types:
            found = True
            break

    _FALKORDB_VECTOR_INDEX_CACHE[key] = found
    return found


def _falkordb_vector_node_search_cypher(
    label: str,
    embedding_attr: str,
    search_vector_param: str,
    use_index: bool,
) -> tuple[str, str]:
    """Build the cypher prefix and node-binding for a node-vector search.

    Returns (prefix, node_var) where:
      - prefix is the Cypher fragment that binds the node variable and a
        `score` variable. With index, it's a CALL ... YIELD; without, it's
        a MATCH plus WITH cosine math.
      - node_var is the variable name the caller's downstream Cypher should
        reference (always 'n' here for parity with the existing code).

    The caller appends WHERE filters and RETURN/ORDER BY/LIMIT as usual.
    The over-fetch parameter `$candidate_k` must be passed by the caller
    when use_index is True.
    """
    if use_index:
        return (
            f"CALL db.idx.vector.queryNodes("
            f"'{label}', '{embedding_attr}', $candidate_k, vecf32({search_vector_param})"
            f") YIELD node, score "
            f"WITH node AS n, score "
        ), "n"
    # Fallback: original cosine math path
    cosine = get_vector_cosine_func_query(
        f"n.{embedding_attr}", search_vector_param, GraphProvider.FALKORDB
    )
    return (
        f"MATCH (n:{label}) "
        f"WITH n, {cosine} AS score "
    ), "n"


def _falkordb_vector_edge_search_cypher(
    relationship_type: str,
    embedding_attr: str,
    search_vector_param: str,
    use_index: bool,
) -> tuple[str, str]:
    """Build the cypher prefix and edge-binding for an edge-vector search.

    Returns (prefix, edge_var). With the index, the procedure binds the
    relationship variable; we then MATCH source and target via the existing
    edge to recover (n)-[e]->(m). Without the index, it's the original
    MATCH-and-cosine path.

    Variable name is 'e' for parity with existing code; source/target are
    'n' and 'm' respectively, also for parity.
    """
    if use_index:
        return (
            f"CALL db.idx.vector.queryRelationships("
            f"'{relationship_type}', '{embedding_attr}', $candidate_k, vecf32({search_vector_param})"
            f") YIELD relationship, score "
            f"MATCH (n:Entity)-[e:{relationship_type}]->(m:Entity) "
            f"WHERE e = relationship "
            f"WITH DISTINCT e, n, m, score "
        ), "e"
    # Fallback
    cosine = get_vector_cosine_func_query(
        f"e.{embedding_attr}", search_vector_param, GraphProvider.FALKORDB
    )
    return (
        f"MATCH (n:Entity)-[e:{relationship_type}]->(m:Entity) "
        f"WITH DISTINCT e, n, m, {cosine} AS score "
    ), "e"



# FalkorDB separator characters that break text into tokens
_SEPARATOR_MAP = str.maketrans(
    {
        ',': ' ',
        '.': ' ',
        '<': ' ',
        '>': ' ',
        '{': ' ',
        '}': ' ',
        '[': ' ',
        ']': ' ',
        '"': ' ',
        "'": ' ',
        ':': ' ',
        ';': ' ',
        '!': ' ',
        '@': ' ',
        '#': ' ',
        '$': ' ',
        '%': ' ',
        '^': ' ',
        '&': ' ',
        '*': ' ',
        '(': ' ',
        ')': ' ',
        '-': ' ',
        '+': ' ',
        '=': ' ',
        '~': ' ',
        '?': ' ',
        '|': ' ',
        '/': ' ',
        '\\': ' ',
    }
)


def _sanitize(query: str) -> str:
    """Replace FalkorDB special characters with whitespace."""
    sanitized = query.translate(_SEPARATOR_MAP)
    return ' '.join(sanitized.split())


def _build_falkor_fulltext_query(
    query: str,
    group_ids: list[str] | None = None,
    max_query_length: int = MAX_QUERY_LENGTH,
) -> str:
    """Build a fulltext query string for FalkorDB using RedisSearch syntax."""
    if group_ids is None or len(group_ids) == 0:
        group_filter = ''
    else:
        escaped_group_ids = [f'"{gid}"' for gid in group_ids]
        group_values = '|'.join(escaped_group_ids)
        group_filter = f'(@group_id:{group_values})'

    sanitized_query = _sanitize(query)

    # Remove stopwords and empty tokens
    query_words = sanitized_query.split()
    filtered_words = [word for word in query_words if word and word.lower() not in STOPWORDS]
    sanitized_query = ' | '.join(filtered_words)

    if len(sanitized_query.split(' ')) + len(group_ids or '') >= max_query_length:
        return ''

    full_query = group_filter + ' (' + sanitized_query + ')'
    return full_query


class FalkorSearchOperations(SearchOperations):
    # --- Node search ---

    async def node_fulltext_search(
        self,
        executor: QueryExecutor,
        query: str,
        search_filter: SearchFilters,
        group_ids: list[str] | None = None,
        limit: int = 10,
    ) -> list[EntityNode]:
        fuzzy_query = _build_falkor_fulltext_query(query, group_ids)
        if fuzzy_query == '':
            return []

        filter_queries, filter_params = node_search_filter_query_constructor(
            search_filter, GraphProvider.FALKORDB
        )

        if group_ids is not None:
            filter_queries.append('n.group_id IN $group_ids')
            filter_params['group_ids'] = group_ids

        filter_query = ''
        if filter_queries:
            filter_query = ' WHERE ' + (' AND '.join(filter_queries))

        cypher = (
            get_nodes_query(
                'node_name_and_summary', '$query', limit=limit, provider=GraphProvider.FALKORDB
            )
            + 'YIELD node AS n, score'
            + filter_query
            + """
            WITH n, score
            ORDER BY score DESC
            LIMIT $limit
            RETURN
            """
            + get_entity_node_return_query(GraphProvider.FALKORDB)
        )

        records, _, _ = await executor.execute_query(
            cypher,
            query=fuzzy_query,
            limit=limit,
            **filter_params,
        )

        return [entity_node_from_record(r) for r in records]

    async def node_similarity_search(
        self,
        executor: QueryExecutor,
        search_vector: list[float],
        search_filter: SearchFilters,
        group_ids: list[str] | None = None,
        limit: int = 10,
        min_score: float = 0.6,
    ) -> list[EntityNode]:
        filter_queries, filter_params = node_search_filter_query_constructor(
            search_filter, GraphProvider.FALKORDB
        )

        if group_ids is not None:
            filter_queries.append('n.group_id IN $group_ids')
            filter_params['group_ids'] = group_ids

        filter_query = ''
        if filter_queries:
            filter_query = ' WHERE ' + (' AND '.join(filter_queries))

        # PATCHED 2026-05-02 (BirdAI vendored patch): use FalkorDB native vector
        # index when available; fall back to interpreted-Cypher cosine math
        # otherwise. The filter clause's position changes between paths
        # (after MATCH for fallback, after YIELD for index path), but the
        # filter expressions themselves are identical because they reference
        # the bound variable `n` either way.
        use_index = await _falkordb_vector_index_exists(
            executor, 'Entity', 'name_embedding', 'NODE'
        )
        prefix, _ = _falkordb_vector_node_search_cypher(
            'Entity', 'name_embedding', '$search_vector', use_index
        )
        where_clauses = []
        if filter_query:
            where_clauses.append(filter_query.replace(' WHERE ', '', 1).strip())
        where_clauses.append('score > $min_score')
        unified_where = ' WHERE ' + ' AND '.join(where_clauses)

        cypher = (
            prefix
            + unified_where
            + """
            RETURN
            """
            + get_entity_node_return_query(GraphProvider.FALKORDB)
            + """
            ORDER BY score DESC
            LIMIT $limit
            """
        )
        params = dict(
            search_vector=search_vector,
            limit=limit,
            min_score=min_score,
            **filter_params,
        )
        if use_index:
            params['candidate_k'] = limit * VECTOR_INDEX_CANDIDATE_MULTIPLIER
        records, _, _ = await executor.execute_query(cypher, **params)

        return [entity_node_from_record(r) for r in records]

    async def node_bfs_search(
        self,
        executor: QueryExecutor,
        origin_uuids: list[str],
        search_filter: SearchFilters,
        max_depth: int,
        group_ids: list[str] | None = None,
        limit: int = 10,
    ) -> list[EntityNode]:
        if not origin_uuids or max_depth < 1:
            return []

        filter_queries, filter_params = node_search_filter_query_constructor(
            search_filter, GraphProvider.FALKORDB
        )

        if group_ids is not None:
            filter_queries.append('n.group_id IN $group_ids')
            filter_queries.append('origin.group_id IN $group_ids')
            filter_params['group_ids'] = group_ids

        filter_query = ''
        if filter_queries:
            filter_query = ' AND ' + (' AND '.join(filter_queries))

        cypher = (
            f"""
            UNWIND $bfs_origin_node_uuids AS origin_uuid
            MATCH (origin {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS*1..{max_depth}]->(n:Entity)
            WHERE n.group_id = origin.group_id
            """
            + filter_query
            + """
            RETURN
            """
            + get_entity_node_return_query(GraphProvider.FALKORDB)
            + """
            LIMIT $limit
            """
        )

        records, _, _ = await executor.execute_query(
            cypher,
            bfs_origin_node_uuids=origin_uuids,
            limit=limit,
            **filter_params,
        )

        return [entity_node_from_record(r) for r in records]

    # --- Edge search ---

    async def edge_fulltext_search(
        self,
        executor: QueryExecutor,
        query: str,
        search_filter: SearchFilters,
        group_ids: list[str] | None = None,
        limit: int = 10,
    ) -> list[EntityEdge]:
        fuzzy_query = _build_falkor_fulltext_query(query, group_ids)
        if fuzzy_query == '':
            return []

        filter_queries, filter_params = edge_search_filter_query_constructor(
            search_filter, GraphProvider.FALKORDB
        )

        if group_ids is not None:
            filter_queries.append('e.group_id IN $group_ids')
            filter_params['group_ids'] = group_ids

        filter_query = ''
        if filter_queries:
            filter_query = ' WHERE ' + (' AND '.join(filter_queries))

        cypher = (
            get_relationships_query(
                'edge_name_and_fact', limit=limit, provider=GraphProvider.FALKORDB
            )
            + """
            YIELD relationship AS rel, score
            MATCH (n:Entity)-[e:RELATES_TO {uuid: rel.uuid}]->(m:Entity)
            """
            + filter_query
            + """
            WITH e, score, n, m
            RETURN
            """
            + get_entity_edge_return_query(GraphProvider.FALKORDB)
            + """
            ORDER BY score DESC
            LIMIT $limit
            """
        )

        records, _, _ = await executor.execute_query(
            cypher,
            query=fuzzy_query,
            limit=limit,
            **filter_params,
        )

        return [entity_edge_from_record(r) for r in records]

    async def edge_similarity_search(
        self,
        executor: QueryExecutor,
        search_vector: list[float],
        source_node_uuid: str | None,
        target_node_uuid: str | None,
        search_filter: SearchFilters,
        group_ids: list[str] | None = None,
        limit: int = 10,
        min_score: float = 0.6,
    ) -> list[EntityEdge]:
        filter_queries, filter_params = edge_search_filter_query_constructor(
            search_filter, GraphProvider.FALKORDB
        )

        if group_ids is not None:
            filter_queries.append('e.group_id IN $group_ids')
            filter_params['group_ids'] = group_ids

            if source_node_uuid is not None:
                filter_params['source_uuid'] = source_node_uuid
                filter_queries.append('n.uuid = $source_uuid')

            if target_node_uuid is not None:
                filter_params['target_uuid'] = target_node_uuid
                filter_queries.append('m.uuid = $target_uuid')

        filter_query = ''
        if filter_queries:
            filter_query = ' WHERE ' + (' AND '.join(filter_queries))

        # PATCHED 2026-05-02 (BirdAI vendored patch): use FalkorDB native vector
        # index on RELATES_TO.fact_embedding when available. The unindexed
        # fallback is the same MATCH-and-cosine math that previously hung
        # for 6+ minutes on a 4,000-entity graph; this is the load-bearing
        # call site that motivated the patch.
        use_index = await _falkordb_vector_index_exists(
            executor, 'RELATES_TO', 'fact_embedding', 'RELATIONSHIP'
        )
        prefix, _ = _falkordb_vector_edge_search_cypher(
            'RELATES_TO', 'fact_embedding', '$search_vector', use_index
        )
        where_clauses = []
        if filter_query:
            where_clauses.append(filter_query.replace(' WHERE ', '', 1).strip())
        where_clauses.append('score > $min_score')
        unified_where = ' WHERE ' + ' AND '.join(where_clauses)

        cypher = (
            prefix
            + unified_where
            + """
            RETURN
            """
            + get_entity_edge_return_query(GraphProvider.FALKORDB)
            + """
            ORDER BY score DESC
            LIMIT $limit
            """
        )
        params = dict(
            search_vector=search_vector,
            limit=limit,
            min_score=min_score,
            **filter_params,
        )
        if use_index:
            params['candidate_k'] = limit * VECTOR_INDEX_CANDIDATE_MULTIPLIER
        records, _, _ = await executor.execute_query(cypher, **params)

        return [entity_edge_from_record(r) for r in records]

    async def edge_bfs_search(
        self,
        executor: QueryExecutor,
        origin_uuids: list[str],
        max_depth: int,
        search_filter: SearchFilters,
        group_ids: list[str] | None = None,
        limit: int = 10,
    ) -> list[EntityEdge]:
        if not origin_uuids:
            return []

        filter_queries, filter_params = edge_search_filter_query_constructor(
            search_filter, GraphProvider.FALKORDB
        )

        if group_ids is not None:
            filter_queries.append('e.group_id IN $group_ids')
            filter_params['group_ids'] = group_ids

        filter_query = ''
        if filter_queries:
            filter_query = ' WHERE ' + (' AND '.join(filter_queries))

        cypher = (
            f"""
            UNWIND $bfs_origin_node_uuids AS origin_uuid
            MATCH path = (origin {{uuid: origin_uuid}})-[:RELATES_TO|MENTIONS*1..{max_depth}]->(:Entity)
            UNWIND relationships(path) AS rel
            MATCH (n:Entity)-[e:RELATES_TO {{uuid: rel.uuid}}]-(m:Entity)
            """
            + filter_query
            + """
            RETURN DISTINCT
            """
            + get_entity_edge_return_query(GraphProvider.FALKORDB)
            + """
            LIMIT $limit
            """
        )

        records, _, _ = await executor.execute_query(
            cypher,
            bfs_origin_node_uuids=origin_uuids,
            depth=max_depth,
            limit=limit,
            **filter_params,
        )

        return [entity_edge_from_record(r) for r in records]

    # --- Episode search ---

    async def episode_fulltext_search(
        self,
        executor: QueryExecutor,
        query: str,
        search_filter: SearchFilters,  # noqa: ARG002
        group_ids: list[str] | None = None,
        limit: int = 10,
    ) -> list[EpisodicNode]:
        fuzzy_query = _build_falkor_fulltext_query(query, group_ids)
        if fuzzy_query == '':
            return []

        filter_params: dict[str, Any] = {}
        group_filter_query = ''
        if group_ids is not None:
            group_filter_query += '\nAND e.group_id IN $group_ids'
            filter_params['group_ids'] = group_ids

        cypher = (
            get_nodes_query(
                'episode_content', '$query', limit=limit, provider=GraphProvider.FALKORDB
            )
            + """
            YIELD node AS episode, score
            MATCH (e:Episodic)
            WHERE e.uuid = episode.uuid
            """
            + group_filter_query
            + """
            RETURN
            """
            + EPISODIC_NODE_RETURN
            + """
            ORDER BY score DESC
            LIMIT $limit
            """
        )

        records, _, _ = await executor.execute_query(
            cypher, query=fuzzy_query, limit=limit, **filter_params
        )

        return [episodic_node_from_record(r) for r in records]

    # --- Community search ---

    async def community_fulltext_search(
        self,
        executor: QueryExecutor,
        query: str,
        group_ids: list[str] | None = None,
        limit: int = 10,
    ) -> list[CommunityNode]:
        fuzzy_query = _build_falkor_fulltext_query(query, group_ids)
        if fuzzy_query == '':
            return []

        filter_params: dict[str, Any] = {}
        group_filter_query = ''
        if group_ids is not None:
            group_filter_query = 'WHERE c.group_id IN $group_ids'
            filter_params['group_ids'] = group_ids

        cypher = (
            get_nodes_query(
                'community_name', '$query', limit=limit, provider=GraphProvider.FALKORDB
            )
            + """
            YIELD node AS c, score
            WITH c, score
            """
            + group_filter_query
            + """
            RETURN
            """
            + COMMUNITY_NODE_RETURN
            + """
            ORDER BY score DESC
            LIMIT $limit
            """
        )

        records, _, _ = await executor.execute_query(
            cypher, query=fuzzy_query, limit=limit, **filter_params
        )

        return [community_node_from_record(r) for r in records]

    async def community_similarity_search(
        self,
        executor: QueryExecutor,
        search_vector: list[float],
        group_ids: list[str] | None = None,
        limit: int = 10,
        min_score: float = 0.6,
    ) -> list[CommunityNode]:
        query_params: dict[str, Any] = {}

        group_filter_query = ''
        if group_ids is not None:
            group_filter_query += ' WHERE c.group_id IN $group_ids'
            query_params['group_ids'] = group_ids

        # PATCHED 2026-05-02 (BirdAI vendored patch): use FalkorDB native vector
        # index on Community.name_embedding when available. Note: the existing
        # filter is built into `group_filter_query` (already prefixed with
        # ' WHERE ' if non-empty) and uses variable `c`. The dispatcher binds
        # the node as `n` for parity with the helper signature, then we
        # re-bind to `c` via WITH so the rest of the query is unchanged.
        use_index = await _falkordb_vector_index_exists(
            executor, 'Community', 'name_embedding', 'NODE'
        )
        prefix, _ = _falkordb_vector_node_search_cypher(
            'Community', 'name_embedding', '$search_vector', use_index
        )
        prefix = prefix + ' WITH n AS c, score '
        where_clauses = []
        if group_filter_query:
            where_clauses.append(group_filter_query.replace(' WHERE ', '', 1).strip())
        where_clauses.append('score > $min_score')
        unified_where = ' WHERE ' + ' AND '.join(where_clauses)

        cypher = (
            prefix
            + unified_where
            + """
            RETURN
            """
            + COMMUNITY_NODE_RETURN
            + """
            ORDER BY score DESC
            LIMIT $limit
            """
        )
        params = dict(
            search_vector=search_vector,
            limit=limit,
            min_score=min_score,
            **query_params,
        )
        if use_index:
            params['candidate_k'] = limit * VECTOR_INDEX_CANDIDATE_MULTIPLIER
        records, _, _ = await executor.execute_query(cypher, **params)

        return [community_node_from_record(r) for r in records]

    # --- Rerankers ---

    async def node_distance_reranker(
        self,
        executor: QueryExecutor,
        node_uuids: list[str],
        center_node_uuid: str,
        min_score: float = 0,
    ) -> list[EntityNode]:
        filtered_uuids = [u for u in node_uuids if u != center_node_uuid]
        scores: dict[str, float] = {center_node_uuid: 0.0}

        cypher = """
        UNWIND $node_uuids AS node_uuid
        MATCH (center:Entity {uuid: $center_uuid})-[:RELATES_TO]-(n:Entity {uuid: node_uuid})
        RETURN 1 AS score, node_uuid AS uuid
        """

        results, _, _ = await executor.execute_query(
            cypher,
            node_uuids=filtered_uuids,
            center_uuid=center_node_uuid,
        )

        for result in results:
            scores[result['uuid']] = result['score']

        for uuid in filtered_uuids:
            if uuid not in scores:
                scores[uuid] = float('inf')

        filtered_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])

        if center_node_uuid in node_uuids:
            scores[center_node_uuid] = 0.1
            filtered_uuids = [center_node_uuid] + filtered_uuids

        reranked_uuids = [u for u in filtered_uuids if (1 / scores[u]) >= min_score]

        if not reranked_uuids:
            return []

        get_query = """
            MATCH (n:Entity)
            WHERE n.uuid IN $uuids
            RETURN
            """ + get_entity_node_return_query(GraphProvider.FALKORDB)

        records, _, _ = await executor.execute_query(get_query, uuids=reranked_uuids)

        node_map = {r['uuid']: entity_node_from_record(r) for r in records}
        return [node_map[u] for u in reranked_uuids if u in node_map]

    async def episode_mentions_reranker(
        self,
        executor: QueryExecutor,
        node_uuids: list[str],
        min_score: float = 0,
    ) -> list[EntityNode]:
        if not node_uuids:
            return []

        scores: dict[str, float] = {}

        results, _, _ = await executor.execute_query(
            """
            UNWIND $node_uuids AS node_uuid
            MATCH (episode:Episodic)-[r:MENTIONS]->(n:Entity {uuid: node_uuid})
            RETURN count(*) AS score, n.uuid AS uuid
            """,
            node_uuids=node_uuids,
        )

        for result in results:
            scores[result['uuid']] = result['score']

        for uuid in node_uuids:
            if uuid not in scores:
                scores[uuid] = float('inf')

        sorted_uuids = list(node_uuids)
        sorted_uuids.sort(key=lambda cur_uuid: scores[cur_uuid])

        reranked_uuids = [u for u in sorted_uuids if scores[u] >= min_score]

        if not reranked_uuids:
            return []

        get_query = """
            MATCH (n:Entity)
            WHERE n.uuid IN $uuids
            RETURN
            """ + get_entity_node_return_query(GraphProvider.FALKORDB)

        records, _, _ = await executor.execute_query(get_query, uuids=reranked_uuids)

        node_map = {r['uuid']: entity_node_from_record(r) for r in records}
        return [node_map[u] for u in reranked_uuids if u in node_map]

    # --- Filter builders ---

    def build_node_search_filters(self, search_filters: SearchFilters) -> Any:
        filter_queries, filter_params = node_search_filter_query_constructor(
            search_filters, GraphProvider.FALKORDB
        )
        return {'filter_queries': filter_queries, 'filter_params': filter_params}

    def build_edge_search_filters(self, search_filters: SearchFilters) -> Any:
        filter_queries, filter_params = edge_search_filter_query_constructor(
            search_filters, GraphProvider.FALKORDB
        )
        return {'filter_queries': filter_queries, 'filter_params': filter_params}

    # --- Fulltext query builder ---

    def build_fulltext_query(
        self,
        query: str,
        group_ids: list[str] | None = None,
        max_query_length: int = MAX_QUERY_LENGTH,
    ) -> str:
        return _build_falkor_fulltext_query(query, group_ids, max_query_length)
