import pytest
import redis
from redis.commands.json.path import Path

from test.testtools import raw_command

json_tests = pytest.importorskip("jsonpath_ng")


def test_arrlen(r: redis.Redis):
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )
    assert (
        r.json().arrlen(
            "arr",
            Path.root_path(),
        )
        == 5
    )
    assert r.json().arrlen("arr") == 5
    assert r.json().arrlen("fake-key") is None

    r.json().set(
        "doc1",
        Path.root_path(),
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )

    assert r.json().arrlen("doc1", "$..a") == [1, 3, None]
    assert r.json().arrlen("doc1", "$.nested1.a") == [3]

    r.json().set(
        "doc2",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", 1, 1, None, "world"]},
            "nested2": {"a": 31},
        },
    )
    assert r.json().arrlen("doc2", "$..a") == [1, 5, None]
    assert r.json().arrlen("doc2", ".nested1.a") == 5
    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )

    # Test multi
    assert r.json().arrlen("doc1", "$..a") == [1, 3, None]
    assert r.json().arrappend("doc1", "$..a", "non", "abba", "stanza") == [
        4,
        6,
        None,
    ]

    r.json().clear("doc1", "$.a")
    assert r.json().arrlen("doc1", "$..a") == [0, 6, None]
    # Test single
    assert r.json().arrlen("doc1", "$.nested1.a") == [6]

    # Test missing key
    with pytest.raises(redis.ResponseError):
        r.json().arrappend("non_existing_doc", "$..a")

    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )
    # Test multi (return result of last path)
    assert r.json().arrlen("doc1", "$..a") == [1, 3, None]
    assert r.json().arrappend("doc1", "..a", "non", "abba", "stanza") == 6

    # Test single
    assert r.json().arrlen("doc1", ".nested1.a") == 6

    # Test missing key
    assert r.json().arrlen("non_existing_doc", "..a") is None


def test_arrappend(r: redis.Redis):
    with pytest.raises(redis.ResponseError):
        r.json().arrappend("non-existing-key", Path.root_path(), 2)

    r.json().set("arr", Path.root_path(), [1])
    assert r.json().arrappend("arr", Path.root_path(), 2) == 2
    assert r.json().arrappend("arr", Path.root_path(), 3, 4) == 4
    assert r.json().arrappend("arr", Path.root_path(), *[5, 6, 7]) == 7
    assert r.json().get("arr") == [1, 2, 3, 4, 5, 6, 7]
    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )
    # Test multi
    assert r.json().arrappend("doc1", "$..a", "bar", "racuda") == [3, 5, None]
    assert r.json().get("doc1", "$") == [
        {
            "a": ["foo", "bar", "racuda"],
            "nested1": {"a": ["hello", None, "world", "bar", "racuda"]},
            "nested2": {"a": 31},
        }
    ]
    assert r.json().arrappend("doc1", "$.nested1.a", "baz") == [6]

    # Test legacy
    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )
    # Test multi (all paths are updated, but return result of last path)
    assert r.json().arrappend("doc1", "..a", "bar", "racuda") == 5

    assert r.json().get("doc1", "$") == [
        {
            "a": ["foo", "bar", "racuda"],
            "nested1": {"a": ["hello", None, "world", "bar", "racuda"]},
            "nested2": {"a": 31},
        }
    ]
    # Test single
    assert r.json().arrappend("doc1", ".nested1.a", "baz") == 6
    assert r.json().get("doc1", "$") == [
        {
            "a": ["foo", "bar", "racuda"],
            "nested1": {"a": ["hello", None, "world", "bar", "racuda", "baz"]},
            "nested2": {"a": 31},
        }
    ]

    # Test missing key
    with pytest.raises(redis.ResponseError):
        r.json().arrappend("non_existing_doc", "$..a")


def test_arrindex(r: redis.Redis):
    r.json().set(
        "foo",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )

    assert r.json().arrindex("foo", Path.root_path(), 1) == 1
    assert r.json().arrindex("foo", Path.root_path(), 1, 2) == -1

    r.json().set(
        "store",
        "$",
        {
            "store": {
                "book": [
                    {
                        "category": "reference",
                        "author": "Nigel Rees",
                        "title": "Sayings of the Century",
                        "price": 8.95,
                        "size": [10, 20, 30, 40],
                    },
                    {
                        "category": "fiction",
                        "author": "Evelyn Waugh",
                        "title": "Sword of Honour",
                        "price": 12.99,
                        "size": [50, 60, 70, 80],
                    },
                    {
                        "category": "fiction",
                        "author": "Herman Melville",
                        "title": "Moby Dick",
                        "isbn": "0-553-21311-3",
                        "price": 8.99,
                        "size": [5, 10, 20, 30],
                    },
                    {
                        "category": "fiction",
                        "author": "J. R. R. Tolkien",
                        "title": "The Lord of the Rings",
                        "isbn": "0-395-19395-8",
                        "price": 22.99,
                        "size": [5, 6, 7, 8],
                    },
                ],
                "bicycle": {"color": "red", "price": 19.95},
            }
        },
    )

    assert r.json().get("store", "$.store.book[?(@.price<10)].size") == [
        [10, 20, 30, 40],
        [5, 10, 20, 30],
    ]
    assert r.json().arrindex("store", "$.store.book[?(@.price<10)].size", "20") == [-1, -1]

    # Test index of int scalar in multi values
    r.json().set(
        "test_num",
        ".",
        [
            {"arr": [0, 1, 3.0, 3, 2, 1, 0, 3]},
            {"nested1_found": {"arr": [5, 4, 3, 2, 1, 0, 1, 2, 3.0, 2, 4, 5]}},
            {"nested2_not_found": {"arr": [2, 4, 6]}},
            {"nested3_scalar": {"arr": "3"}},
            [
                {"nested41_not_arr": {"arr_renamed": [1, 2, 3]}},
                {"nested42_empty_arr": {"arr": []}},
            ],
        ],
    )

    assert r.json().get("test_num", "$..arr") == [
        [0, 1, 3.0, 3, 2, 1, 0, 3],
        [5, 4, 3, 2, 1, 0, 1, 2, 3.0, 2, 4, 5],
        [2, 4, 6],
        "3",
        [],
    ]

    assert r.json().arrindex("test_num", "$..nonexistingpath", 3) == []
    assert r.json().arrindex("test_num", "$..arr", 3) == [3, 2, -1, None, -1]

    # Test index of double scalar in multi values
    assert r.json().arrindex("test_num", "$..arr", 3.0) == [2, 8, -1, None, -1]

    # Test index of string scalar in multi values
    r.json().set(
        "test_string",
        ".",
        [
            {"arr": ["bazzz", "bar", 2, "baz", 2, "ba", "baz", 3]},
            {"nested1_found": {"arr": [None, "baz2", "buzz", 2, 1, 0, 1, "2", "baz", 2, 4, 5]}},
            {"nested2_not_found": {"arr": ["baz2", 4, 6]}},
            {"nested3_scalar": {"arr": "3"}},
            [
                {"nested41_arr": {"arr_renamed": [1, "baz", 3]}},
                {"nested42_empty_arr": {"arr": []}},
            ],
        ],
    )
    assert r.json().get("test_string", "$..arr") == [
        ["bazzz", "bar", 2, "baz", 2, "ba", "baz", 3],
        [None, "baz2", "buzz", 2, 1, 0, 1, "2", "baz", 2, 4, 5],
        ["baz2", 4, 6],
        "3",
        [],
    ]

    assert r.json().arrindex("test_string", "$..arr", "baz") == [
        3,
        8,
        -1,
        None,
        -1,
    ]

    assert r.json().arrindex("test_string", "$..arr", "baz", 2) == [
        3,
        8,
        -1,
        None,
        -1,
    ]
    assert r.json().arrindex("test_string", "$..arr", "baz", 4) == [
        6,
        8,
        -1,
        None,
        -1,
    ]
    assert r.json().arrindex("test_string", "$..arr", "baz", -5) == [
        3,
        8,
        -1,
        None,
        -1,
    ]
    assert r.json().arrindex("test_string", "$..arr", "baz", 4, 7) == [
        6,
        -1,
        -1,
        None,
        -1,
    ]
    assert r.json().arrindex("test_string", "$..arr", "baz", 4, -1) == [
        6,
        8,
        -1,
        None,
        -1,
    ]
    assert r.json().arrindex("test_string", "$..arr", "baz", 4, 0) == [
        6,
        8,
        -1,
        None,
        -1,
    ]
    assert r.json().arrindex("test_string", "$..arr", "5", 7, -1) == [
        -1,
        -1,
        -1,
        None,
        -1,
    ]
    assert r.json().arrindex("test_string", "$..arr", "5", 7, 0) == [
        -1,
        -1,
        -1,
        None,
        -1,
    ]

    # Test index of None scalar in multi values
    r.json().set(
        "test_None",
        ".",
        [
            {"arr": ["bazzz", "None", 2, None, 2, "ba", "baz", 3]},
            {"nested1_found": {"arr": ["zaz", "baz2", "buzz", 2, 1, 0, 1, "2", None, 2, 4, 5]}},
            {"nested2_not_found": {"arr": ["None", 4, 6]}},
            {"nested3_scalar": {"arr": None}},
            [
                {"nested41_arr": {"arr_renamed": [1, None, 3]}},
                {"nested42_empty_arr": {"arr": []}},
            ],
        ],
    )
    assert r.json().get("test_None", "$..arr") == [
        ["bazzz", "None", 2, None, 2, "ba", "baz", 3],
        ["zaz", "baz2", "buzz", 2, 1, 0, 1, "2", None, 2, 4, 5],
        ["None", 4, 6],
        None,
        [],
    ]

    # Test with none-scalar value
    # assert r.json().arrindex("test_None", "$..nested42_empty_arr.arr", {"arr": []}) == [-1]

    # Test legacy (path begins with dot)
    # Test index of int scalar in single value
    assert r.json().arrindex("test_num", ".[0].arr", 3) == 3
    assert r.json().arrindex("test_num", ".[0].arr", 9) == -1

    with pytest.raises(redis.ResponseError):
        r.json().arrindex("test_num", ".[0].arr_not", 3)
    # Test index of string scalar in single value
    assert r.json().arrindex("test_string", ".[0].arr", "baz") == 3
    assert r.json().arrindex("test_string", ".[0].arr", "faz") == -1
    # Test index of None scalar in single value
    assert r.json().arrindex("test_None", ".[0].arr", "None") == 1
    assert r.json().arrindex("test_None", "..nested2_not_found.arr", "None") == 0


def test_arrinsert(r: redis.Redis):
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 4],
    )

    assert r.json().arrinsert("arr", Path.root_path(), 1, *[1, 2, 3]) == 5
    assert r.json().get("arr") == [0, 1, 2, 3, 4]

    # test prepends
    r.json().set("val2", Path.root_path(), [5, 6, 7, 8, 9])
    assert r.json().arrinsert("val2", Path.root_path(), 0, ["some", "thing"]) == 6
    assert r.json().get("val2") == [["some", "thing"], 5, 6, 7, 8, 9]
    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )
    # Test multi
    assert r.json().arrinsert("doc1", "$..a", "1", "bar", "racuda") == [3, 5, None]

    assert r.json().get("doc1", "$") == [
        {
            "a": ["foo", "bar", "racuda"],
            "nested1": {"a": ["hello", "bar", "racuda", None, "world"]},
            "nested2": {"a": 31},
        }
    ]
    # Test single
    assert r.json().arrinsert("doc1", "$.nested1.a", -2, "baz") == [6]
    assert r.json().get("doc1", "$") == [
        {
            "a": ["foo", "bar", "racuda"],
            "nested1": {"a": ["hello", "bar", "racuda", "baz", None, "world"]},
            "nested2": {"a": 31},
        }
    ]

    # Test missing key
    with pytest.raises(redis.ResponseError):
        r.json().arrappend("non_existing_doc", "$..a")


def test_arrpop(r: redis.Redis):
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )
    assert raw_command(r, "json.arrpop", "arr") == b"4"

    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )
    assert r.json().arrpop("arr", Path.root_path(), 4) == 4
    assert r.json().arrpop("arr", Path.root_path(), -1) == 3
    assert r.json().arrpop("arr", Path.root_path()) == 2
    assert r.json().arrpop("arr", Path.root_path(), 0) == 0
    assert r.json().get("arr") == [1]

    # test out of bounds
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )
    assert r.json().arrpop("arr", Path.root_path(), 99) == 4

    # none test
    r.json().set(
        "arr",
        Path.root_path(),
        [],
    )
    assert r.json().arrpop("arr") is None

    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )

    # # Test legacy
    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )
    # Test multi (all paths are updated, but return result of last path)
    assert r.json().arrpop("doc1", "..a", "1") is None
    assert r.json().get("doc1", "$") == [{"a": [], "nested1": {"a": ["hello", "world"]}, "nested2": {"a": 31}}]

    # # Test missing key
    with pytest.raises(redis.ResponseError):
        r.json().arrpop("non_existing_doc", "..a")


def test_arrtrim(r: redis.Redis):
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )

    assert r.json().arrtrim("arr", Path.root_path(), 1, 3) == 3
    assert r.json().get("arr") == [1, 2, 3]

    # <0 test, should be 0 equivalent
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )
    assert r.json().arrtrim("arr", Path.root_path(), -1, 3) == 0

    # testing stop > end
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )
    assert r.json().arrtrim("arr", Path.root_path(), 3, 99) == 2

    # start > array size and stop
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )
    assert r.json().arrtrim("arr", Path.root_path(), 9, 1) == 0

    # all larger
    r.json().set(
        "arr",
        Path.root_path(),
        [0, 1, 2, 3, 4],
    )
    assert r.json().arrtrim("arr", Path.root_path(), 9, 11) == 0

    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )
    # Test multi
    assert r.json().arrtrim("doc1", "$..a", "1", -1) == [0, 2, None]
    assert r.json().get("doc1", "$") == [{"a": [], "nested1": {"a": [None, "world"]}, "nested2": {"a": 31}}]

    r.json().set("doc1", "$", {"a": [], "nested1": {"a": [None, "world"]}, "nested2": {"a": 31}})
    assert r.json().arrtrim("doc1", "$..a", "1", "1") == [0, 1, None]
    assert r.json().get("doc1", "$") == [{"a": [], "nested1": {"a": ["world"]}, "nested2": {"a": 31}}]
    # Test single
    assert r.json().arrtrim("doc1", "$.nested1.a", 1, 0) == [0]
    assert r.json().get("doc1", "$") == [{"a": [], "nested1": {"a": []}, "nested2": {"a": 31}}]

    # Test missing key
    with pytest.raises(redis.ResponseError):
        r.json().arrtrim("non_existing_doc", "..a", "0", 1)

    # Test legacy
    r.json().set(
        "doc1",
        "$",
        {
            "a": ["foo"],
            "nested1": {"a": ["hello", None, "world"]},
            "nested2": {"a": 31},
        },
    )

    # Test multi (all paths are updated, but return result of last path)
    assert r.json().arrtrim("doc1", "..a", "1", "-1") == 2

    # Test single
    assert r.json().arrtrim("doc1", ".nested1.a", "1", "1") == 1
    assert r.json().get("doc1", "$") == [{"a": [], "nested1": {"a": ["world"]}, "nested2": {"a": 31}}]

    # Test missing key
    with pytest.raises(redis.ResponseError):
        r.json().arrtrim("non_existing_doc", "..a", 1, 1)