"""Tests for COUNT(*), COUNT(DISTINCT x), and DISTINCT inside aggregates."""

import networkx as nx

from nxcypher import NXCypher


class TestCountStar:
    """Tests for COUNT(*) — count all matching rows."""

    def test_count_star_basic(self):
        """COUNT(*) counts all nodes."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")
        host.add_node("c", name="Charlie")

        res = NXCypher(host).run("MATCH (n) RETURN COUNT(*)")
        assert res["COUNT(*)"] == [3]

    def test_count_star_with_edges(self):
        """COUNT(*) counts all matched edge patterns."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")
        host.add_node("c", name="Charlie")
        host.add_edge("a", "b", __labels__={"KNOWS"})
        host.add_edge("a", "c", __labels__={"KNOWS"})
        host.add_edge("b", "c", __labels__={"KNOWS"})

        res = NXCypher(host).run("""
            MATCH (a)-[:KNOWS]->(b)
            RETURN COUNT(*)
        """)
        assert res["COUNT(*)"] == [3]

    def test_count_star_with_grouping(self):
        """COUNT(*) with a grouping key."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")
        host.add_node("c", name="Charlie")
        host.add_edge("a", "b", __labels__={"KNOWS"})
        host.add_edge("a", "c", __labels__={"KNOWS"})
        host.add_edge("b", "c", __labels__={"KNOWS"})

        res = NXCypher(host).run("""
            MATCH (x)-[:KNOWS]->(y)
            RETURN x.name, COUNT(*)
        """)
        # Alice->Bob, Alice->Charlie = 2; Bob->Charlie = 1
        alice_idx = res["x.name"].index("Alice")
        bob_idx = res["x.name"].index("Bob")
        assert res["COUNT(*)"][alice_idx] == 2
        assert res["COUNT(*)"][bob_idx] == 1

    def test_count_star_no_matches(self):
        """COUNT(*) with no matches returns [0]."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")

        res = NXCypher(host).run("""
            MATCH (n {name: "Nonexistent"})
            RETURN COUNT(*)
        """)
        assert res["COUNT(*)"] == [0]

    def test_count_star_with_alias(self):
        """COUNT(*) with AS alias."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")

        res = NXCypher(host).run("MATCH (n) RETURN COUNT(*) AS total")
        assert res["total"] == [2]

    def test_count_star_order_by(self):
        """COUNT(*) used in ORDER BY."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")
        host.add_node("c", name="Charlie")
        host.add_edge("a", "b", __labels__={"KNOWS"})
        host.add_edge("a", "c", __labels__={"KNOWS"})
        host.add_edge("b", "c", __labels__={"KNOWS"})

        res = NXCypher(host).run("""
            MATCH (x)-[:KNOWS]->(y)
            RETURN x.name, COUNT(*)
            ORDER BY COUNT(*) DESC
        """)
        assert res["x.name"][0] == "Alice"
        assert res["COUNT(*)"][0] == 2

    def test_count_star_with_other_aggregates(self):
        """COUNT(*) alongside other aggregates."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")
        host.add_node("c", name="Charlie")
        host.add_edge("a", "b", __labels__={"KNOWS"}, weight=10)
        host.add_edge("a", "c", __labels__={"KNOWS"}, weight=20)

        res = NXCypher(host).run("""
            MATCH (x)-[r:KNOWS]->(y)
            RETURN x.name, COUNT(*), SUM(r.weight)
        """)
        assert res["x.name"] == ["Alice"]
        assert res["COUNT(*)"] == [2]
        assert res["SUM(r.weight)"] == [30]

    def test_count_star_in_with_clause(self):
        """COUNT(*) in WITH clause."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")
        host.add_node("c", name="Charlie")
        host.add_edge("a", "b", __labels__={"KNOWS"})
        host.add_edge("a", "c", __labels__={"KNOWS"})
        host.add_edge("b", "c", __labels__={"KNOWS"})

        res = NXCypher(host).run("""
            MATCH (x)-[:KNOWS]->(y)
            WITH x.name AS person, COUNT(*) AS cnt
            RETURN person, cnt
        """)
        alice_idx = res["person"].index("Alice")
        bob_idx = res["person"].index("Bob")
        assert res["cnt"][alice_idx] == 2
        assert res["cnt"][bob_idx] == 1


class TestCountDistinct:
    """Tests for COUNT(DISTINCT x)."""

    def test_count_distinct_basic(self):
        """COUNT(DISTINCT x) counts unique values."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice", city="NYC")
        host.add_node("b", name="Bob", city="NYC")
        host.add_node("c", name="Charlie", city="LA")

        res = NXCypher(host).run("MATCH (n) RETURN COUNT(DISTINCT n.city)")
        assert res["COUNT(DISTINCT n.city)"] == [2]

    def test_count_distinct_with_grouping(self):
        """COUNT(DISTINCT x) with grouping key."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")
        host.add_node("c", name="Charlie")
        host.add_node("d", name="Dave")
        host.add_edge("a", "b", __labels__={"KNOWS"}, kind="friend")
        host.add_edge("a", "c", __labels__={"KNOWS"}, kind="friend")
        host.add_edge("a", "d", __labels__={"KNOWS"}, kind="colleague")
        host.add_edge("b", "c", __labels__={"KNOWS"}, kind="friend")

        res = NXCypher(host).run("""
            MATCH (x)-[r:KNOWS]->(y)
            RETURN x.name, COUNT(DISTINCT r.kind)
        """)
        alice_idx = res["x.name"].index("Alice")
        bob_idx = res["x.name"].index("Bob")
        # Alice has 2 distinct kinds: friend, colleague
        assert res["COUNT(DISTINCT r.kind)"][alice_idx] == 2
        # Bob has 1 distinct kind: friend
        assert res["COUNT(DISTINCT r.kind)"][bob_idx] == 1

    def test_count_distinct_with_alias(self):
        """COUNT(DISTINCT x) with alias."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice", city="NYC")
        host.add_node("b", name="Bob", city="NYC")
        host.add_node("c", name="Charlie", city="LA")

        res = NXCypher(host).run(
            "MATCH (n) RETURN COUNT(DISTINCT n.city) AS unique_cities"
        )
        assert res["unique_cities"] == [2]

    def test_count_distinct_vs_count(self):
        """COUNT(DISTINCT x) vs COUNT(x) return different values."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice", color="red")
        host.add_node("b", name="Bob", color="red")
        host.add_node("c", name="Charlie", color="blue")

        res_distinct = NXCypher(host).run("MATCH (n) RETURN COUNT(DISTINCT n.color)")
        res_count = NXCypher(host).run("MATCH (n) RETURN COUNT(n.color)")

        assert res_distinct["COUNT(DISTINCT n.color)"] == [2]
        assert res_count["COUNT(n.color)"] == [3]


class TestCollectDistinct:
    """Tests for COLLECT(DISTINCT x)."""

    def test_collect_distinct_basic(self):
        """COLLECT(DISTINCT x) returns unique values."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice", city="NYC")
        host.add_node("b", name="Bob", city="NYC")
        host.add_node("c", name="Charlie", city="LA")

        res = NXCypher(host).run("MATCH (n) RETURN COLLECT(DISTINCT n.city)")
        collected = res["COLLECT(DISTINCT n.city)"][0]
        assert set(collected) == {"NYC", "LA"}
        assert len(collected) == 2

    def test_collect_distinct_with_grouping(self):
        """COLLECT(DISTINCT x) with grouping key."""
        host = nx.DiGraph()
        host.add_node("a", name="Alice")
        host.add_node("b", name="Bob")
        host.add_node("c", name="Charlie")
        host.add_edge("a", "b", __labels__={"TAG"}, tag="python")
        host.add_edge("a", "c", __labels__={"TAG"}, tag="python")
        host.add_edge("a", "b", __labels__={"TAG"}, tag="java")

        res = NXCypher(host).run("""
            MATCH (x)-[r:TAG]->(y)
            RETURN x.name, COLLECT(DISTINCT r.tag)
        """)
        assert res["x.name"] == ["Alice"]
        collected = res["COLLECT(DISTINCT r.tag)"][0]
        assert set(collected) == {"python", "java"}
        assert len(collected) == 2
