"""Tests for UNION and UNION ALL clauses."""

import pytest
import networkx as nx

from nxcypher import NXCypher


class TestUnionAll:
    """Tests for UNION ALL - keeps all rows including duplicates."""

    def test_union_all_basic(self):
        """UNION ALL combines results from two queries."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", type="Person")
        host.add_node("Bob", name="Bob", type="Person")
        host.add_node("Charlie", name="Charlie", type="Robot")

        qry = """
        MATCH (n {type: "Person"})
        RETURN n.name
        UNION ALL
        MATCH (n {type: "Robot"})
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        assert len(res["n.name"]) == 3
        assert "Alice" in res["n.name"]
        assert "Bob" in res["n.name"]
        assert "Charlie" in res["n.name"]

    def test_union_all_with_duplicates(self):
        """UNION ALL keeps duplicate rows."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", age=30)
        host.add_node("Bob", name="Bob", age=30)

        qry = """
        MATCH (n)
        RETURN n.age
        UNION ALL
        MATCH (n)
        RETURN n.age
        """
        res = NXCypher(host).run(qry)

        # Should have 4 rows (2 from each query, with duplicates)
        assert len(res["n.age"]) == 4
        assert res["n.age"].count(30) == 4

    def test_union_all_different_columns(self):
        """UNION ALL with queries returning different column aliases."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")

        # Both queries return same column name
        qry = """
        MATCH (n {name: "Alice"})
        RETURN n.name
        UNION ALL
        MATCH (n {name: "Bob"})
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        assert len(res["n.name"]) == 2
        assert "Alice" in res["n.name"]
        assert "Bob" in res["n.name"]

    def test_union_all_multiple(self):
        """Multiple UNION ALL clauses."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", role="admin")
        host.add_node("Bob", name="Bob", role="user")
        host.add_node("Charlie", name="Charlie", role="guest")

        qry = """
        MATCH (n {role: "admin"})
        RETURN n.name
        UNION ALL
        MATCH (n {role: "user"})
        RETURN n.name
        UNION ALL
        MATCH (n {role: "guest"})
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        assert len(res["n.name"]) == 3
        assert "Alice" in res["n.name"]
        assert "Bob" in res["n.name"]
        assert "Charlie" in res["n.name"]


class TestUnionDistinct:
    """Tests for UNION - removes duplicate rows."""

    def test_union_basic(self):
        """UNION combines results and removes duplicates."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", type="Person")
        host.add_node("Bob", name="Bob", type="Person")

        qry = """
        MATCH (n)
        RETURN n.name
        UNION
        MATCH (n)
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        # Should have 2 rows (duplicates removed)
        assert len(res["n.name"]) == 2
        assert "Alice" in res["n.name"]
        assert "Bob" in res["n.name"]

    def test_union_removes_duplicates(self):
        """UNION removes duplicate rows across queries."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", age=30)
        host.add_node("Bob", name="Bob", age=25)

        qry = """
        MATCH (n {name: "Alice"})
        RETURN n.age
        UNION
        MATCH (n)
        RETURN n.age
        """
        res = NXCypher(host).run(qry)

        # Alice's age appears in both queries, should be deduplicated
        assert len(res["n.age"]) == 2
        assert 30 in res["n.age"]
        assert 25 in res["n.age"]

    def test_union_multiple(self):
        """Multiple UNION clauses."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", role="admin")
        host.add_node("Bob", name="Bob", role="user")
        host.add_node("Charlie", name="Charlie", role="user")

        qry = """
        MATCH (n {role: "admin"})
        RETURN n.name
        UNION
        MATCH (n {role: "user"})
        RETURN n.name
        UNION
        MATCH (n {role: "user"})
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        # Should have 3 unique names (Bob and Charlie from user queries deduplicated)
        assert len(res["n.name"]) == 3
        assert "Alice" in res["n.name"]
        assert "Bob" in res["n.name"]
        assert "Charlie" in res["n.name"]


class TestUnionMixed:
    """Tests for mixed UNION and UNION ALL."""

    def test_union_and_union_all_mixed(self):
        """Mix of UNION and UNION ALL in same query."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")

        qry = """
        MATCH (n {name: "Alice"})
        RETURN n.name
        UNION ALL
        MATCH (n {name: "Alice"})
        RETURN n.name
        UNION
        MATCH (n {name: "Alice"})
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        # First UNION ALL keeps duplicates (2 Alices)
        # Second UNION removes duplicates from the third query
        # So we have: Alice, Alice from first two, then third Alice is duplicate
        assert len(res["n.name"]) == 2  # UNION at end removes the duplicate


class TestUnionWithRelationships:
    """Tests for UNION with relationship patterns."""

    def test_union_with_relationship_match(self):
        """UNION with queries that match relationships."""
        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("Bob", "Charlie", __labels__={"WORKS_WITH"})

        qry = """
        MATCH (a)-[:KNOWS]->(b)
        RETURN b.name
        UNION ALL
        MATCH (a)-[:WORKS_WITH]->(b)
        RETURN b.name
        """
        res = NXCypher(host).run(qry)

        assert len(res["b.name"]) == 2
        assert "Bob" in res["b.name"]
        assert "Charlie" in res["b.name"]

    def test_union_different_relationship_types(self):
        """UNION combining different relationship patterns."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")
        host.add_node("Bob", name="Bob")
        host.add_node("Work", name="Work")
        host.add_edge("Alice", "Bob", __labels__={"KNOWS"})
        host.add_edge("Alice", "Work", __labels__={"WORKS_AT"})

        qry = """
        MATCH (a {name: "Alice"})-[:KNOWS]->(target)
        RETURN target.name
        UNION ALL
        MATCH (a {name: "Alice"})-[:WORKS_AT]->(target)
        RETURN target.name
        """
        res = NXCypher(host).run(qry)

        assert len(res["target.name"]) == 2
        assert "Bob" in res["target.name"]
        assert "Work" in res["target.name"]


class TestUnionEdgeCases:
    """Edge cases for UNION operations."""

    def test_union_empty_first_query(self):
        """UNION where first query returns no results."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", type="Person")

        qry = """
        MATCH (n {type: "Robot"})
        RETURN n.name
        UNION ALL
        MATCH (n {type: "Person"})
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        assert len(res["n.name"]) == 1
        assert res["n.name"] == ["Alice"]

    def test_union_empty_second_query(self):
        """UNION where second query returns no results."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", type="Person")

        qry = """
        MATCH (n {type: "Person"})
        RETURN n.name
        UNION ALL
        MATCH (n {type: "Robot"})
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        assert len(res["n.name"]) == 1
        assert res["n.name"] == ["Alice"]

    def test_union_multiple_columns(self):
        """UNION with multiple return columns."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", age=30)
        host.add_node("Bob", name="Bob", age=25)

        qry = """
        MATCH (n {name: "Alice"})
        RETURN n.name, n.age
        UNION ALL
        MATCH (n {name: "Bob"})
        RETURN n.name, n.age
        """
        res = NXCypher(host).run(qry)

        assert len(res["n.name"]) == 2
        assert len(res["n.age"]) == 2

        # Check both rows are present
        rows = list(zip(res["n.name"], res["n.age"]))
        assert ("Alice", 30) in rows
        assert ("Bob", 25) in rows

    def test_union_with_where_clause(self):
        """UNION with WHERE clauses in both queries."""
        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)

        qry = """
        MATCH (n)
        WHERE n.age > 28
        RETURN n.name
        UNION ALL
        MATCH (n)
        WHERE n.age < 28
        RETURN n.name
        """
        res = NXCypher(host).run(qry)

        assert len(res["n.name"]) == 3
        assert "Alice" in res["n.name"]
        assert "Bob" in res["n.name"]
        assert "Charlie" in res["n.name"]
