"""Surreal-native Graphiti search interface tests."""

from __future__ import annotations

from datetime import UTC, datetime

import pytest

from sibyl_core.backends.surreal import SurrealDriver
from sibyl_core.backends.surreal.schema import EMBEDDING_DIM
from sibyl_core.graph.search_interface import SurrealSearchInterface
from sibyl_core.graph.surreal.compat.models import (
    CommunityNode,
    EntityEdge,
    EntityNode,
    EpisodeType,
    EpisodicEdge,
    EpisodicNode,
    HasEpisodeEdge,
    SagaNode,
)
from sibyl_core.graph.surreal.compat.search_filters import (
    ComparisonOperator,
    DateFilter,
    SearchFilters,
)


def _entity(uuid: str, group_id: str, *, name: str, summary: str) -> EntityNode:
    return EntityNode(
        uuid=uuid,
        name=name,
        name_embedding=[0.1] * EMBEDDING_DIM,
        group_id=group_id,
        labels=["Entity", "Pattern"],
        created_at=datetime.now(UTC),
        summary=summary,
        attributes={},
    )


def _edge(uuid: str, group_id: str, source: str, target: str, *, fact: str) -> EntityEdge:
    now = datetime.now(UTC)
    return EntityEdge(
        uuid=uuid,
        source_node_uuid=source,
        target_node_uuid=target,
        fact=fact,
        fact_embedding=[0.1] * EMBEDDING_DIM,
        name="RELATES_TO",
        group_id=group_id,
        episodes=[],
        created_at=now,
        expired_at=None,
        valid_at=now,
        invalid_at=None,
        attributes={},
    )


def _episode(uuid: str, group_id: str, *, content: str) -> EpisodicNode:
    now = datetime.now(UTC)
    return EpisodicNode(
        uuid=uuid,
        name="Surreal episode",
        content=content,
        source=EpisodeType.text,
        source_description="",
        group_id=group_id,
        created_at=now,
        valid_at=now,
        entity_edges=[],
    )


def _mention(uuid: str, group_id: str, source: str, target: str) -> EpisodicEdge:
    return EpisodicEdge(
        uuid=uuid,
        group_id=group_id,
        source_node_uuid=source,
        target_node_uuid=target,
        created_at=datetime.now(UTC),
    )


def _community(uuid: str, group_id: str, *, name: str, summary: str) -> CommunityNode:
    return CommunityNode(
        uuid=uuid,
        name=name,
        group_id=group_id,
        summary=summary,
        name_embedding=[0.2] * EMBEDDING_DIM,
        created_at=datetime.now(UTC),
    )


class TestSurrealSearchInterfaceIntegration:
    @pytest.mark.asyncio
    async def test_node_fulltext_and_similarity_search(self, surreal_schema: SurrealDriver) -> None:
        gid = surreal_schema.group_id
        interface = SurrealSearchInterface()
        await surreal_schema.entity_node_ops.save(
            surreal_schema,
            _entity("node-search", gid, name="Surreality Prism", summary="northstar memory"),
        )

        fulltext = await interface.node_fulltext_search(
            surreal_schema,
            "surreality",
            SearchFilters(node_labels=["Pattern"]),
            [gid],
            5,
        )
        similarity = await interface.node_similarity_search(
            surreal_schema,
            [0.1] * EMBEDDING_DIM,
            SearchFilters(node_labels=["Pattern"]),
            [gid],
            5,
            0.0,
        )

        assert [node.uuid for node in fulltext] == ["node-search"]
        assert [node.uuid for node in similarity] == ["node-search"]

    @pytest.mark.asyncio
    async def test_edge_fulltext_and_similarity_search(self, surreal_schema: SurrealDriver) -> None:
        gid = surreal_schema.group_id
        interface = SurrealSearchInterface()
        await surreal_schema.entity_node_ops.save(
            surreal_schema,
            _entity("source-node", gid, name="Source", summary="source"),
        )
        await surreal_schema.entity_node_ops.save(
            surreal_schema,
            _entity("target-node", gid, name="Target", summary="target"),
        )
        await surreal_schema.entity_edge_ops.save(
            surreal_schema,
            _edge(
                "edge-search",
                gid,
                "source-node",
                "target-node",
                fact="Surreality links coding memory",
            ),
        )

        fulltext = await interface.edge_fulltext_search(
            surreal_schema,
            "surreality",
            SearchFilters(edge_uuids=["edge-search"], node_labels=["Pattern"]),
            [gid],
            5,
        )
        similarity = await interface.edge_similarity_search(
            surreal_schema,
            [0.1] * EMBEDDING_DIM,
            "source-node",
            "target-node",
            SearchFilters(
                edge_uuids=["edge-search"],
                valid_at=[
                    [
                        DateFilter(
                            date=datetime(2100, 1, 1, tzinfo=UTC),
                            comparison_operator=ComparisonOperator.less_than_equal,
                        )
                    ]
                ],
            ),
            [gid],
            5,
            0.0,
        )
        filtered_out = await interface.edge_similarity_search(
            surreal_schema,
            [0.1] * EMBEDDING_DIM,
            "source-node",
            "target-node",
            SearchFilters(
                edge_uuids=["edge-search"],
                valid_at=[
                    [
                        DateFilter(
                            date=datetime(2020, 1, 1, tzinfo=UTC),
                            comparison_operator=ComparisonOperator.less_than,
                        )
                    ]
                ],
            ),
            [gid],
            5,
            0.0,
        )

        assert [edge.uuid for edge in fulltext] == ["edge-search"]
        assert [edge.uuid for edge in similarity] == ["edge-search"]
        assert filtered_out == []

    @pytest.mark.asyncio
    async def test_episode_fulltext_search(self, surreal_schema: SurrealDriver) -> None:
        gid = surreal_schema.group_id
        interface = SurrealSearchInterface()
        await surreal_schema.episode_node_ops.save(
            surreal_schema,
            _episode("episode-search", gid, content="Surreality captures raw memory"),
        )

        results = await interface.episode_fulltext_search(
            surreal_schema,
            "surreality",
            SearchFilters(),
            [gid],
            5,
        )

        assert [episode.uuid for episode in results] == ["episode-search"]

    @pytest.mark.asyncio
    async def test_community_fulltext_similarity_and_embedding_load(
        self, surreal_schema: SurrealDriver
    ) -> None:
        gid = surreal_schema.group_id
        interface = SurrealSearchInterface()
        community = _community(
            "community-search",
            gid,
            name="Surreality Guild",
            summary="native community search",
        )
        await surreal_schema.community_node_ops.save(surreal_schema, community)

        fulltext = await interface.community_fulltext_search(
            surreal_schema,
            "surreality",
            [gid],
            5,
        )
        similarity = await interface.community_similarity_search(
            surreal_schema,
            [0.2] * EMBEDDING_DIM,
            [gid],
            5,
            0.0,
        )
        embeddings = await interface.get_embeddings_for_communities(
            surreal_schema,
            [community],
        )
        blank_fulltext = await interface.community_fulltext_search(
            surreal_schema,
            "   ",
            [gid],
            5,
        )
        blank_similarity = await interface.community_similarity_search(
            surreal_schema,
            [],
            [gid],
            5,
            0.0,
        )
        wrong_group = await interface.community_fulltext_search(
            surreal_schema,
            "surreality",
            ["other-group"],
            5,
        )

        assert [community.uuid for community in fulltext] == ["community-search"]
        assert [community.uuid for community in similarity] == ["community-search"]
        assert embeddings["community-search"] == [0.2] * EMBEDDING_DIM
        assert blank_fulltext == []
        assert blank_similarity == []
        assert wrong_group == []

    @pytest.mark.asyncio
    async def test_local_model_methods_use_surreal_graph_operations_interface(
        self, surreal_schema: SurrealDriver
    ) -> None:
        gid = surreal_schema.group_id
        source = _entity("graphiti-source", gid, name="Graphiti Source", summary="source")
        target = _entity("graphiti-target", gid, name="Graphiti Target", summary="target")
        edge = _edge(
            "graphiti-edge",
            gid,
            "graphiti-source",
            "graphiti-target",
            fact="Graphiti native operation edge",
        )
        episode = _episode(
            "graphiti-episode",
            gid,
            content="Graphiti native operation episode",
        )

        await source.save(surreal_schema)
        await target.save(surreal_schema)
        await edge.save(surreal_schema)
        await episode.save(surreal_schema)

        loaded_source = await EntityNode.get_by_uuid(surreal_schema, "graphiti-source")
        loaded_edge = await EntityEdge.get_by_uuid(surreal_schema, "graphiti-edge")
        recent_episodes = await surreal_schema.graph_operations_interface.retrieve_episodes(
            surreal_schema,
            datetime.now(UTC),
            1,
            [gid],
        )
        node_embeddings = await surreal_schema.graph_operations_interface.node_load_embeddings_bulk(
            surreal_schema,
            [loaded_source],
        )
        edge_embeddings = await surreal_schema.graph_operations_interface.edge_load_embeddings_bulk(
            surreal_schema,
            [loaded_edge],
        )

        assert loaded_source.uuid == "graphiti-source"
        assert loaded_edge.uuid == "graphiti-edge"
        assert [episode.uuid for episode in recent_episodes] == ["graphiti-episode"]
        assert node_embeddings["graphiti-source"] == [0.1] * EMBEDDING_DIM
        assert edge_embeddings["graphiti-edge"] == [0.1] * EMBEDDING_DIM

    @pytest.mark.asyncio
    async def test_graphiti_bulk_dict_payloads_use_surreal_graph_operations_interface(
        self, surreal_schema: SurrealDriver
    ) -> None:
        gid = surreal_schema.group_id
        now = datetime.now(UTC)
        interface = surreal_schema.graph_operations_interface

        await interface.node_save_bulk(
            None,
            surreal_schema,
            None,
            [
                {
                    "uuid": "bulk-source",
                    "name": "Bulk Source",
                    "group_id": gid,
                    "labels": ["Entity", "Tool"],
                    "summary": "source",
                    "created_at": now,
                    "name_embedding": [0.1] * EMBEDDING_DIM,
                    "project_id": "project-bulk",
                },
                {
                    "uuid": "bulk-target",
                    "name": "Bulk Target",
                    "group_id": gid,
                    "labels": ["Entity", "Tool"],
                    "summary": "target",
                    "created_at": now,
                    "name_embedding": [0.1] * EMBEDDING_DIM,
                },
            ],
        )
        await interface.edge_save_bulk(
            None,
            surreal_schema,
            None,
            [
                {
                    "uuid": "bulk-edge",
                    "source_node_uuid": "bulk-source",
                    "target_node_uuid": "bulk-target",
                    "name": "RELATES_TO",
                    "fact": "Bulk source relates to bulk target",
                    "fact_embedding": [0.1] * EMBEDDING_DIM,
                    "group_id": gid,
                    "episodes": [],
                    "created_at": now,
                    "expired_at": None,
                    "valid_at": now,
                    "invalid_at": None,
                    "confidence": 0.9,
                }
            ],
        )

        source = await EntityNode.get_by_uuid(surreal_schema, "bulk-source")
        edge = await EntityEdge.get_by_uuid(surreal_schema, "bulk-edge")

        assert source.attributes["project_id"] == "project-bulk"
        assert edge.attributes["confidence"] == 0.9

    @pytest.mark.asyncio
    async def test_saga_helpers_use_surreal_graph_operations_interface(
        self, surreal_schema: SurrealDriver
    ) -> None:
        gid = surreal_schema.group_id
        interface = surreal_schema.graph_operations_interface
        old_at = datetime(2026, 1, 1, tzinfo=UTC)
        current_at = datetime(2026, 1, 2, tzinfo=UTC)

        await surreal_schema.execute_query(
            "CREATE saga SET uuid = 'saga-interface', name = 'daily', "
            "group_id = $gid, created_at = $created_at;",
            gid=gid,
            created_at=old_at,
        )
        for uuid, content, created_at in (
            ("saga-old-episode", "old content", old_at),
            ("saga-current-episode", "current content", current_at),
        ):
            await surreal_schema.episode_node_ops.save(
                surreal_schema,
                _episode(uuid, gid, content=content),
            )
            await surreal_schema.execute_query(
                "UPDATE episode SET created_at = $created_at, valid_at = $created_at "
                "WHERE uuid = $uuid;",
                uuid=uuid,
                created_at=created_at,
            )
            await surreal_schema.has_episode_edge_ops.save(
                surreal_schema,
                HasEpisodeEdge(
                    uuid=f"has-{uuid}",
                    group_id=gid,
                    source_node_uuid="saga-interface",
                    target_node_uuid=uuid,
                    created_at=created_at,
                ),
            )

        previous_uuid = await interface.saga_get_previous_episode_uuid(
            surreal_schema,
            "saga-interface",
            "saga-current-episode",
        )
        latest_contents = await interface.saga_get_episode_contents(
            surreal_schema,
            "saga-interface",
            limit=2,
        )
        new_contents = await interface.saga_get_episode_contents(
            surreal_schema,
            "saga-interface",
            since=datetime(2026, 1, 1, 12, tzinfo=UTC),
        )

        assert previous_uuid == "saga-old-episode"
        assert latest_contents == ["old content", "current content"]
        assert new_contents == ["current content"]

    @pytest.mark.asyncio
    async def test_saga_episode_bulk_ops_use_surreal_adapter(
        self, surreal_schema: SurrealDriver
    ) -> None:
        gid = surreal_schema.group_id
        now = datetime(2026, 1, 1, tzinfo=UTC)
        interface = surreal_schema.graph_operations_interface

        await interface.saga_node_save(
            SagaNode(uuid="daily-saga", name="daily", group_id=gid, created_at=now),
            surreal_schema,
        )
        await interface.episodic_node_save_bulk(
            None,
            surreal_schema,
            None,
            [
                {
                    "uuid": "bulk-episode",
                    "name": "Bulk episode",
                    "content": "Bliss is testing Surreal episode compatibility.",
                    "source": EpisodeType.message,
                    "source_description": "test",
                    "group_id": gid,
                    "created_at": now,
                    "valid_at": now,
                    "entity_edges": [],
                }
            ],
        )
        await interface.has_episode_edge_save_bulk(
            None,
            surreal_schema,
            None,
            [
                HasEpisodeEdge(
                    uuid="has-bulk-episode",
                    group_id=gid,
                    source_node_uuid="daily-saga",
                    target_node_uuid="bulk-episode",
                    created_at=now,
                )
            ],
        )

        episodes = await surreal_schema.execute_query("SELECT uuid, name FROM episode;")
        sagas = await surreal_schema.execute_query("SELECT uuid, name FROM saga;")
        saga_edges = await surreal_schema.execute_query("SELECT uuid FROM has_episode;")

        assert [episode["name"] for episode in episodes] == ["Bulk episode"]
        assert [episode["uuid"] for episode in episodes] == ["bulk-episode"]
        assert [saga["name"] for saga in sagas] == ["daily"]
        assert len(saga_edges) == 1

    @pytest.mark.asyncio
    async def test_episode_related_payload_deletion_uses_surreal_native_ops(
        self, surreal_schema: SurrealDriver
    ) -> None:
        gid = surreal_schema.group_id
        interface = surreal_schema.graph_operations_interface
        await surreal_schema.episode_node_ops.save(
            surreal_schema,
            _episode("remove-episode", gid, content="removal regression"),
        )
        await surreal_schema.execute_query(
            "UPDATE episode SET entity_edges = ['remove-edge'] WHERE uuid = 'remove-episode';"
        )
        for node in (
            _entity("remove-node", gid, name="Remove", summary="mentioned once"),
            _entity("keep-node", gid, name="Keep", summary="edge target"),
        ):
            await surreal_schema.entity_node_ops.save(surreal_schema, node)
        await surreal_schema.entity_edge_ops.save(
            surreal_schema,
            _edge(
                "remove-edge",
                gid,
                "remove-node",
                "keep-node",
                fact="episode-created edge",
            ),
        )
        await surreal_schema.execute_query(
            "UPDATE relates_to SET episodes = ['remove-episode'] WHERE uuid = 'remove-edge';"
        )
        await surreal_schema.episodic_edge_ops.save(
            surreal_schema,
            _mention("remove-mention", gid, "remove-episode", "remove-node"),
        )

        await interface.episodic_edge_delete_by_uuids(
            None,
            surreal_schema,
            ["remove-mention"],
            gid,
        )
        await interface.edge_delete_by_uuids(
            None,
            surreal_schema,
            ["remove-edge"],
            gid,
        )
        await interface.node_delete(
            await EntityNode.get_by_uuid(surreal_schema, "remove-node"),
            surreal_schema,
        )
        await interface.episodic_node_delete(
            await EpisodicNode.get_by_uuid(surreal_schema, "remove-episode"),
            surreal_schema,
        )

        episodes = await surreal_schema.execute_query(
            "SELECT uuid FROM episode WHERE uuid = 'remove-episode';"
        )
        removed_nodes = await surreal_schema.execute_query(
            "SELECT uuid FROM entity WHERE uuid = 'remove-node';"
        )
        kept_nodes = await surreal_schema.execute_query(
            "SELECT uuid FROM entity WHERE uuid = 'keep-node';"
        )
        edges = await surreal_schema.execute_query(
            "SELECT uuid FROM relates_to WHERE uuid = 'remove-edge';"
        )
        mentions = await surreal_schema.execute_query(
            "SELECT uuid FROM mentions WHERE uuid = 'remove-mention';"
        )

        assert episodes == []
        assert removed_nodes == []
        assert [node["uuid"] for node in kept_nodes] == ["keep-node"]
        assert edges == []
        assert mentions == []

    @pytest.mark.asyncio
    async def test_bfs_and_rerankers_use_native_surreal_queries(
        self, surreal_schema: SurrealDriver
    ) -> None:
        gid = surreal_schema.group_id
        interface = SurrealSearchInterface()
        for node in (
            _entity("bfs-origin", gid, name="Origin", summary="origin"),
            _entity("bfs-mid", gid, name="Middle", summary="middle"),
            _entity("bfs-target", gid, name="Target", summary="target"),
            _entity("bfs-cross-group", "other-group", name="Cross", summary="cross"),
        ):
            await surreal_schema.entity_node_ops.save(surreal_schema, node)
        for episode in (
            _episode("bfs-episode-a", gid, content="mentions middle"),
            _episode("bfs-episode-b", gid, content="mentions middle again"),
            _episode("bfs-episode-c", gid, content="mentions target"),
        ):
            await surreal_schema.episode_node_ops.save(surreal_schema, episode)
        for edge in (
            _edge("bfs-edge-1", gid, "bfs-origin", "bfs-mid", fact="origin to middle"),
            _edge("bfs-edge-2", gid, "bfs-mid", "bfs-target", fact="middle to target"),
            _edge(
                "bfs-cross-edge",
                gid,
                "bfs-origin",
                "bfs-cross-group",
                fact="cross group leak candidate",
            ),
        ):
            await surreal_schema.entity_edge_ops.save(surreal_schema, edge)
        for mention in (
            _mention("mention-a", gid, "bfs-episode-a", "bfs-mid"),
            _mention("mention-b", gid, "bfs-episode-b", "bfs-mid"),
            _mention("mention-c", gid, "bfs-episode-c", "bfs-target"),
            _mention("mention-cross", gid, "bfs-episode-a", "bfs-cross-group"),
        ):
            await surreal_schema.episodic_edge_ops.save(surreal_schema, mention)

        nodes_from_entity = await interface.node_bfs_search(
            surreal_schema,
            ["bfs-origin"],
            SearchFilters(node_labels=["Pattern"]),
            2,
            [gid],
            5,
        )
        nodes_from_episode = await interface.node_bfs_search(
            surreal_schema,
            ["bfs-episode-a"],
            SearchFilters(node_labels=["Pattern"]),
            2,
            [gid],
            5,
        )
        edges_from_entity = await interface.edge_bfs_search(
            surreal_schema,
            ["bfs-origin"],
            2,
            SearchFilters(),
            [gid],
            5,
        )
        edges_from_episode = await interface.edge_bfs_search(
            surreal_schema,
            ["bfs-episode-a"],
            2,
            SearchFilters(),
            [gid],
            5,
        )
        distance_uuids, distance_scores = await interface.node_distance_reranker(
            surreal_schema,
            ["bfs-target", "bfs-mid", "bfs-origin"],
            "bfs-origin",
        )
        mention_uuids, mention_scores = await interface.episode_mentions_reranker(
            surreal_schema,
            [["bfs-target", "bfs-mid"], ["bfs-mid"]],
            1,
        )
        scoped_distance_uuids, scoped_distance_scores = await interface.node_distance_reranker(
            surreal_schema,
            ["bfs-mid", "bfs-cross-group"],
            "bfs-origin",
            1,
        )
        scoped_mention_uuids, scoped_mention_scores = await interface.episode_mentions_reranker(
            surreal_schema,
            [["bfs-mid", "bfs-cross-group"]],
            1,
        )

        assert [node.uuid for node in nodes_from_entity] == ["bfs-mid", "bfs-target"]
        assert [node.uuid for node in nodes_from_episode] == ["bfs-mid", "bfs-target"]
        assert [edge.uuid for edge in edges_from_entity] == ["bfs-edge-1", "bfs-edge-2"]
        assert [edge.uuid for edge in edges_from_episode] == ["bfs-edge-2"]
        assert distance_uuids == ["bfs-origin", "bfs-mid", "bfs-target"]
        assert distance_scores == [0.1, 1.0, 0.0]
        assert mention_uuids == ["bfs-mid", "bfs-target"]
        assert mention_scores == [2.0, 1.0]
        assert scoped_distance_uuids == ["bfs-mid"]
        assert scoped_distance_scores == [1.0]
        assert scoped_mention_uuids == ["bfs-mid"]
        assert scoped_mention_scores == [2.0]
