"""Tests for list comprehension expressions."""

import pytest
import networkx as nx

from nxcypher import NXCypher


def get_list_comp_result(res):
    """Helper to get the list comprehension result from the result dict."""
    for key in res.keys():
        if key.startswith("LISTCOMP_"):
            return res[key]
    return None


class TestListComprehensionBasic:
    """Basic tests for list comprehensions."""

    def test_list_comp_identity(self):
        """List comprehension that just returns elements."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3] | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [1, 2, 3]

    def test_list_comp_multiply(self):
        """List comprehension with multiplication transformation."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3] | x * 2]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [2, 4, 6]

    def test_list_comp_add(self):
        """List comprehension with addition transformation."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3] | x + 10]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [11, 12, 13]

    def test_list_comp_subtract(self):
        """List comprehension with subtraction transformation."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [10, 20, 30] | x - 5]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [5, 15, 25]

    def test_list_comp_divide(self):
        """List comprehension with division transformation."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [10, 20, 30] | x / 2]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [5.0, 10.0, 15.0]


class TestListComprehensionFilter:
    """Tests for list comprehensions with WHERE filter."""

    def test_list_comp_filter_greater_than(self):
        """List comprehension with > filter."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3, 4, 5] WHERE x > 2 | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [3, 4, 5]

    def test_list_comp_filter_less_than(self):
        """List comprehension with < filter."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3, 4, 5] WHERE x < 4 | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [1, 2, 3]

    def test_list_comp_filter_gte(self):
        """List comprehension with >= filter."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3, 4, 5] WHERE x >= 3 | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [3, 4, 5]

    def test_list_comp_filter_lte(self):
        """List comprehension with <= filter."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3, 4, 5] WHERE x <= 3 | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [1, 2, 3]

    def test_list_comp_filter_eq(self):
        """List comprehension with = filter."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3, 4, 5] WHERE x = 3 | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [3]

    def test_list_comp_filter_neq(self):
        """List comprehension with <> filter."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3, 4, 5] WHERE x <> 3 | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [1, 2, 4, 5]


class TestListComprehensionFilterAndTransform:
    """Tests for list comprehensions with both filter and transform."""

    def test_list_comp_filter_and_multiply(self):
        """List comprehension with filter and multiply."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3, 4, 5] WHERE x > 2 | x * 2]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [6, 8, 10]

    def test_list_comp_filter_and_add(self):
        """List comprehension with filter and add."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3, 4, 5] WHERE x < 4 | x + 100]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [101, 102, 103]


class TestListComprehensionWithNodeProperty:
    """Tests for list comprehensions using node properties as source."""

    def test_list_comp_with_node_property_list(self):
        """List comprehension using a node property as source."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", scores=[10, 20, 30, 40])

        qry = """
        MATCH (n {name: "Alice"})
        RETURN [x IN n.scores | x * 2]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [20, 40, 60, 80]

    def test_list_comp_with_node_property_filter(self):
        """List comprehension with filter on node property list."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", scores=[10, 20, 30, 40, 50])

        qry = """
        MATCH (n {name: "Alice"})
        RETURN [x IN n.scores WHERE x > 25 | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [30, 40, 50]


class TestListComprehensionEdgeCases:
    """Edge case tests for list comprehensions."""

    def test_list_comp_empty_list(self):
        """List comprehension on empty list."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", items=[])

        qry = """
        MATCH (n {name: "Alice"})
        RETURN [x IN n.items | x * 2]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == []

    def test_list_comp_filter_removes_all(self):
        """List comprehension where filter removes all elements."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [1, 2, 3] WHERE x > 100 | x]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == []

    def test_list_comp_single_element(self):
        """List comprehension on single-element list."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice")

        qry = """
        MATCH (n)
        RETURN [x IN [5] | x * 3]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)
        assert result[0] == [15]

    def test_list_comp_multiple_rows(self):
        """List comprehension applied to multiple matched nodes."""
        host = nx.MultiDiGraph()
        host.add_node("Alice", name="Alice", values=[1, 2])
        host.add_node("Bob", name="Bob", values=[10, 20])

        qry = """
        MATCH (n)
        RETURN n.name, [x IN n.values | x * 2]
        """
        res = NXCypher(host).run(qry)
        result = get_list_comp_result(res)

        # Find results by name
        names = res["n.name"]
        alice_idx = names.index("Alice")
        bob_idx = names.index("Bob")

        assert result[alice_idx] == [2, 4]
        assert result[bob_idx] == [20, 40]
