"""Tests for COLLECT() aggregate function."""

import pytest
import networkx as nx

from nxcypher import NXCypher


class TestCollectBasic:
    """Basic COLLECT() functionality tests."""

    def test_collect_single_property(self):
        """COLLECT basic property values into a list."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", age=30)
        host.add_node("Bob", name="Bob", age=25)
        host.add_node("Charlie", name="Charlie", age=35)
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})
        host.add_edge("Alice", "Charlie", __labels__={"KNOWS"})

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

    def test_collect_multiple_groups(self):
        """COLLECT with multiple grouping keys."""
        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)
        RETURN p.name, COLLECT(item.name)
        """
        res = NXCypher(host).run(qry)
        # Alice owns X, Y; Bob owns Z
        alice_idx = res["p.name"].index("Alice")
        bob_idx = res["p.name"].index("Bob")
        assert set(res["COLLECT(item.name)"][alice_idx]) == {"X", "Y"}
        assert res["COLLECT(item.name)"][bob_idx] == ["Z"]

    def test_collect_empty_result(self):
        """COLLECT with no matching patterns returns empty result."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

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

    def test_collect_with_none_values(self):
        """COLLECT handles None values in collected data."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")
        host.add_node("Charlie")  # No name property
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})
        host.add_edge("Alice", "Charlie", __labels__={"KNOWS"})

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

    def test_collect_node_attrs(self):
        """COLLECT node returns list of attribute dicts."""
        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)
        RETURN a.name, COLLECT(b)
        """
        res = NXCypher(host).run(qry)
        assert res["a.name"] == ["Alice"]
        # Collecting a node returns list of attribute dicts
        collected = res["COLLECT(b)"][0]
        assert len(collected) == 2
        names = {item["name"] for item in collected}
        assert names == {"Bob", "Charlie"}


class TestCollectWithOtherAggregates:
    """Tests combining COLLECT with other aggregate functions."""

    def test_collect_with_count(self):
        """Use COLLECT alongside COUNT in same query."""
        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)
        RETURN a.name, COUNT(b), COLLECT(b.name)
        """
        res = NXCypher(host).run(qry)
        assert res["a.name"] == ["Alice"]
        assert res["COUNT(b)"] == [2]
        assert set(res["COLLECT(b.name)"][0]) == {"Bob", "Charlie"}


class TestCollectEdgeAttributes:
    """Tests for collecting edge attributes."""

    def test_collect_edge_property(self):
        """COLLECT edge property values."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")
        host.add_edge("Alice", "Bob", __labels__={"PAID"}, amount=100)
        host.add_edge("Alice", "Bob", __labels__={"PAID"}, amount=50)

        qry = """
        MATCH (a)-[r:PAID]->(b)
        RETURN a.name, b.name, COLLECT(r.amount)
        """
        res = NXCypher(host).run(qry)
        assert res["a.name"] == ["Alice"]
        assert res["b.name"] == ["Bob"]
        assert set(res["COLLECT(r.amount)"][0]) == {100, 50}


class TestCollectWithOrderBy:
    """Tests for COLLECT with ORDER BY."""

    def test_collect_order_by_group_key(self):
        """COLLECT results ordered by grouping key."""
        host = nx.MultiDiGraph()
        host.add_node("Bob", name="Bob")
        host.add_node("Alice", name="Alice")
        host.add_node("X", name="X")
        host.add_node("Y", name="Y")
        host.add_edge("Bob", "X", __labels__={"OWNS"})
        host.add_edge("Alice", "Y", __labels__={"OWNS"})

        qry = """
        MATCH (p)-[:OWNS]->(item)
        RETURN p.name, COLLECT(item.name)
        ORDER BY p.name
        """
        res = NXCypher(host).run(qry)
        assert res["p.name"] == ["Alice", "Bob"]
        assert res["COLLECT(item.name)"][0] == ["Y"]
        assert res["COLLECT(item.name)"][1] == ["X"]


class TestCollectWithLimit:
    """Tests for COLLECT with LIMIT."""

    def test_collect_with_limit(self):
        """COLLECT respects LIMIT on 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)
        RETURN p.name, COLLECT(item.name)
        LIMIT 1
        """
        res = NXCypher(host).run(qry)
        assert len(res["p.name"]) == 1
