"""Tests for WITH clause."""

import pytest
import networkx as nx

from nxcypher import NXCypher


class TestWithBasic:
    """Basic WITH clause functionality tests."""

    def test_with_passthrough(self):
        """WITH passes variables through unchanged."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", age=30)
        host.add_node("Bob", name="Bob", age=25)
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})

        qry = """
        MATCH (a)-[:KNOWS]->(b)
        WITH a, b
        RETURN a.name, b.name
        """
        res = NXCypher(host).run(qry)
        assert res["a.name"] == ["Alice"]
        assert res["b.name"] == ["Bob"]

    def test_with_alias(self):
        """WITH can rename variables."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", age=30)
        host.add_node("Bob", name="Bob", age=25)
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})

        qry = """
        MATCH (a)-[:KNOWS]->(b)
        WITH a.name AS person, b.name AS friend
        RETURN person, friend
        """
        res = NXCypher(host).run(qry)
        assert res["person"] == ["Alice"]
        assert res["friend"] == ["Bob"]

    def test_with_count_aggregation(self):
        """WITH with COUNT aggregation."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")
        host.add_node("Charlie", name="Charlie")
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})
        host.add_edge("Alice", "Charlie", __labels__={"KNOWS"})

        qry = """
        MATCH (a)-[:KNOWS]->(b)
        WITH a.name AS person, COUNT(b) AS num_friends
        RETURN person, num_friends
        """
        res = NXCypher(host).run(qry)
        assert res["person"] == ["Alice"]
        assert res["num_friends"] == [2]

    def test_with_collect_aggregation(self):
        """WITH with COLLECT aggregation."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")
        host.add_node("Charlie", name="Charlie")
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})
        host.add_edge("Alice", "Charlie", __labels__={"KNOWS"})

        qry = """
        MATCH (a)-[:KNOWS]->(b)
        WITH a.name AS person, COLLECT(b.name) AS friends
        RETURN person, friends
        """
        res = NXCypher(host).run(qry)
        assert res["person"] == ["Alice"]
        assert set(res["friends"][0]) == {"Bob", "Charlie"}


class TestWithFiltering:
    """Tests for WITH clause filtering."""

    def test_with_filters_results(self):
        """WITH limits which variables pass through."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", age=30)
        host.add_node("Bob", name="Bob", age=25)
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})

        qry = """
        MATCH (a)-[:KNOWS]->(b)
        WITH a
        RETURN a.name
        """
        res = NXCypher(host).run(qry)
        assert res["a.name"] == ["Alice"]


class TestWithMultipleGroups:
    """Tests for WITH with multiple grouping scenarios."""

    def test_with_multiple_groups(self):
        """WITH handles multiple groups."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")
        host.add_node("X", name="X")
        host.add_node("Y", name="Y")
        host.add_node("Z", name="Z")
        host.add_edge("Alice", "X", __labels__={"OWNS"})
        host.add_edge("Alice", "Y", __labels__={"OWNS"})
        host.add_edge("Bob", "Z", __labels__={"OWNS"})

        qry = """
        MATCH (p)-[:OWNS]->(item)
        WITH p.name AS person, COUNT(item) AS item_count
        RETURN person, item_count
        """
        res = NXCypher(host).run(qry)
        # Alice owns 2 items, Bob owns 1
        alice_idx = res["person"].index("Alice")
        bob_idx = res["person"].index("Bob")
        assert res["item_count"][alice_idx] == 2
        assert res["item_count"][bob_idx] == 1


class TestWithDistinct:
    """Tests for WITH DISTINCT."""

    def test_with_distinct(self):
        """WITH DISTINCT removes duplicates."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")
        host.add_node("Charlie", name="Charlie")
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})
        host.add_edge("Alice", "Charlie", __labels__={"KNOWS"})

        qry = """
        MATCH (a)-[:KNOWS]->(b)
        WITH DISTINCT a.name AS person
        RETURN person
        """
        res = NXCypher(host).run(qry)
        assert res["person"] == ["Alice"]
