# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
"""Tests for JSON graph serialization/deserialization roundtrips."""

from __future__ import annotations

import json
from typing import Any, Callable

import pytest
import tvm_ffi
import tvm_ffi.testing
from tvm_ffi.serialization import from_json_graph_str, to_json_graph_str


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------
def _roundtrip(obj: Any) -> Any:
    """Serialize then deserialize and return the result."""
    return from_json_graph_str(to_json_graph_str(obj))


def _assert_roundtrip_eq(obj: Any, cmp: Callable[..., Any] | None = None) -> None:
    """Assert that roundtrip preserves the value."""
    result = _roundtrip(obj)
    if cmp is not None:
        cmp(result)
    elif isinstance(obj, float):
        assert isinstance(result, float)
        assert result == pytest.approx(obj)
    elif obj is None:
        assert result is None
    else:
        _assert_any_equal(result, obj)


def _assert_any_equal(a: Any, b: Any) -> None:
    """Recursively compare two tvm_ffi values for equality."""
    if isinstance(b, tvm_ffi.Array):
        assert len(a) == len(b)
        for x, y in zip(a, b):
            _assert_any_equal(x, y)
    elif isinstance(b, (tvm_ffi.Map, tvm_ffi.Dict)):
        assert len(a) == len(b)
        for k in b:
            _assert_any_equal(a[k], b[k])
    elif isinstance(b, tvm_ffi.Shape):
        assert list(a) == list(b)
    elif isinstance(b, str):
        # tvm_ffi String inherits from str
        assert str(a) == str(b)
    else:
        assert a == b


# ---------------------------------------------------------------------------
# Primitive types
# ---------------------------------------------------------------------------
class TestNone:
    """Roundtrip tests for None."""

    def test_none(self) -> None:
        """None roundtrips to None."""
        assert _roundtrip(None) is None


class TestBool:
    """Roundtrip tests for bool values."""

    def test_true(self) -> None:
        """True roundtrips to True."""
        assert _roundtrip(True) is True

    def test_false(self) -> None:
        """False roundtrips to False."""
        assert _roundtrip(False) is False


class TestInt:
    """Roundtrip tests for integer values."""

    def test_zero(self) -> None:
        """Zero roundtrips correctly."""
        result = _roundtrip(0)
        assert result == 0

    def test_positive(self) -> None:
        """Positive int roundtrips correctly."""
        result = _roundtrip(42)
        assert result == 42

    def test_negative(self) -> None:
        """Negative int roundtrips correctly."""
        result = _roundtrip(-1)
        assert result == -1

    def test_large_positive(self) -> None:
        """Large positive int roundtrips correctly."""
        result = _roundtrip(10**15)
        assert result == 10**15

    def test_large_negative(self) -> None:
        """Large negative int roundtrips correctly."""
        result = _roundtrip(-(10**15))
        assert result == -(10**15)


class TestFloat:
    """Roundtrip tests for float values."""

    def test_zero(self) -> None:
        """Float zero roundtrips correctly."""
        result = _roundtrip(0.0)
        assert result == 0.0

    def test_positive(self) -> None:
        """Positive float roundtrips correctly."""
        result = _roundtrip(3.14159)
        assert result == pytest.approx(3.14159)

    def test_negative(self) -> None:
        """Negative float roundtrips correctly."""
        result = _roundtrip(-2.718)
        assert result == pytest.approx(-2.718)

    def test_very_small(self) -> None:
        """Very small float roundtrips correctly."""
        result = _roundtrip(1e-300)
        assert result == pytest.approx(1e-300)

    def test_very_large(self) -> None:
        """Very large float roundtrips correctly."""
        result = _roundtrip(1e300)
        assert result == pytest.approx(1e300)


# ---------------------------------------------------------------------------
# String types
# ---------------------------------------------------------------------------
class TestString:
    """Roundtrip tests for ffi.String values."""

    def test_empty(self) -> None:
        """Empty string roundtrips correctly."""
        _assert_roundtrip_eq(tvm_ffi.convert(""))

    def test_short(self) -> None:
        """Short string roundtrips correctly."""
        _assert_roundtrip_eq(tvm_ffi.convert("hello"))

    def test_long(self) -> None:
        """Long string roundtrips correctly."""
        _assert_roundtrip_eq(tvm_ffi.convert("x" * 1000))

    def test_special_chars(self) -> None:
        """String with special characters roundtrips correctly."""
        _assert_roundtrip_eq(tvm_ffi.convert('hello\nworld\t"quotes"'))

    def test_unicode(self) -> None:
        """Unicode string roundtrips correctly."""
        _assert_roundtrip_eq(tvm_ffi.convert("hello 世界"))


# ---------------------------------------------------------------------------
# DataType
# ---------------------------------------------------------------------------
class TestDataType:
    """Roundtrip tests for DLDataType values."""

    def test_int32(self) -> None:
        """int32 dtype roundtrips correctly."""
        s = to_json_graph_str(tvm_ffi.dtype("int32"))
        result = from_json_graph_str(s)
        assert str(result) == "int32"

    def test_float64(self) -> None:
        """float64 dtype roundtrips correctly."""
        s = to_json_graph_str(tvm_ffi.dtype("float64"))
        result = from_json_graph_str(s)
        assert str(result) == "float64"

    def test_bool(self) -> None:
        """Bool dtype roundtrips correctly."""
        s = to_json_graph_str(tvm_ffi.dtype("bool"))
        result = from_json_graph_str(s)
        assert str(result) == "bool"

    def test_vector(self) -> None:
        """Vector dtype roundtrips correctly."""
        s = to_json_graph_str(tvm_ffi.dtype("float32x4"))
        result = from_json_graph_str(s)
        assert str(result) == "float32x4"


# ---------------------------------------------------------------------------
# Device
# ---------------------------------------------------------------------------
class TestDevice:
    """Roundtrip tests for DLDevice values."""

    def test_cpu(self) -> None:
        """CPU device roundtrips correctly."""
        s = to_json_graph_str(tvm_ffi.Device("cpu", 0))
        result = from_json_graph_str(s)
        assert result.dlpack_device_type() == tvm_ffi.Device("cpu", 0).dlpack_device_type()
        assert result.index == 0

    def test_cuda(self) -> None:
        """CUDA device roundtrips correctly."""
        s = to_json_graph_str(tvm_ffi.Device("cuda", 1))
        result = from_json_graph_str(s)
        assert result.dlpack_device_type() == tvm_ffi.Device("cuda", 1).dlpack_device_type()
        assert result.index == 1


# ---------------------------------------------------------------------------
# Containers
# ---------------------------------------------------------------------------
class TestArray:
    """Roundtrip tests for ffi.Array containers."""

    def test_empty(self) -> None:
        """Empty array roundtrips correctly."""
        arr = tvm_ffi.convert([])
        _assert_roundtrip_eq(arr)

    def test_single_element(self) -> None:
        """Single-element array roundtrips correctly."""
        arr = tvm_ffi.convert([42])
        result = _roundtrip(arr)
        assert len(result) == 1
        assert result[0] == 42

    def test_multiple_elements(self) -> None:
        """Multi-element array roundtrips correctly."""
        arr = tvm_ffi.convert([1, 2, 3])
        result = _roundtrip(arr)
        assert len(result) == 3
        assert list(result) == [1, 2, 3]

    def test_mixed_types(self) -> None:
        """Array with mixed types roundtrips correctly."""
        arr = tvm_ffi.convert([42, "hello", True, None])
        result = _roundtrip(arr)
        assert len(result) == 4
        assert result[0] == 42
        assert result[1] == "hello"
        assert result[2] is True
        assert result[3] is None

    def test_nested_arrays(self) -> None:
        """Nested arrays roundtrip correctly."""
        inner1 = tvm_ffi.convert([1, 2])
        inner2 = tvm_ffi.convert([3])
        outer = tvm_ffi.convert([inner1, inner2])
        result = _roundtrip(outer)
        assert len(result) == 2
        assert list(result[0]) == [1, 2]
        assert list(result[1]) == [3]

    def test_duplicated_elements(self) -> None:
        """Array with duplicated elements roundtrips correctly."""
        arr = tvm_ffi.convert([42, 42, 42])
        result = _roundtrip(arr)
        assert len(result) == 3
        assert all(x == 42 for x in result)


class TestMap:
    """Roundtrip tests for ffi.Map containers."""

    def test_empty(self) -> None:
        """Empty map roundtrips correctly."""
        m = tvm_ffi.convert({})
        _assert_roundtrip_eq(m)

    def test_single_entry(self) -> None:
        """Single-entry map roundtrips correctly."""
        m = tvm_ffi.convert({"key": 42})
        result = _roundtrip(m)
        assert len(result) == 1
        assert result["key"] == 42

    def test_multiple_entries(self) -> None:
        """Multi-entry map roundtrips correctly."""
        m = tvm_ffi.convert({"a": 1, "b": 2, "c": 3})
        result = _roundtrip(m)
        assert len(result) == 3
        assert result["a"] == 1
        assert result["b"] == 2
        assert result["c"] == 3

    def test_mixed_value_types(self) -> None:
        """Map with mixed value types roundtrips correctly."""
        m = tvm_ffi.convert({"int": 42, "str": "hello", "bool": True, "none": None})
        result = _roundtrip(m)
        assert result["int"] == 42
        assert result["str"] == "hello"
        assert result["bool"] is True
        assert result["none"] is None

    def test_nested_map(self) -> None:
        """Nested maps roundtrip correctly."""
        inner = tvm_ffi.convert({"x": 1})
        outer = tvm_ffi.convert({"inner": inner})
        result = _roundtrip(outer)
        assert result["inner"]["x"] == 1

    def test_map_with_array_value(self) -> None:
        """Map with array values roundtrips correctly."""
        arr = tvm_ffi.convert([10, 20])
        m = tvm_ffi.convert({"nums": arr})
        result = _roundtrip(m)
        assert list(result["nums"]) == [10, 20]


class TestDict:
    """Roundtrip tests for ffi.Dict containers."""

    def test_empty(self) -> None:
        """Empty dict roundtrips correctly."""
        d = tvm_ffi.Dict({})
        _assert_roundtrip_eq(d)

    def test_single_entry(self) -> None:
        """Single-entry dict roundtrips correctly."""
        d = tvm_ffi.Dict({"key": 42})
        result = _roundtrip(d)
        assert len(result) == 1
        assert result["key"] == 42

    def test_multiple_entries(self) -> None:
        """Multi-entry dict roundtrips correctly."""
        d = tvm_ffi.Dict({"a": 1, "b": 2, "c": 3})
        result = _roundtrip(d)
        assert len(result) == 3
        assert result["a"] == 1
        assert result["b"] == 2
        assert result["c"] == 3

    def test_mixed_value_types(self) -> None:
        """Dict with mixed value types roundtrips correctly."""
        d = tvm_ffi.Dict({"int": 42, "str": "hello", "bool": True, "none": None})
        result = _roundtrip(d)
        assert result["int"] == 42
        assert result["str"] == "hello"
        assert result["bool"] is True
        assert result["none"] is None

    def test_nested_dict(self) -> None:
        """Nested dicts roundtrip correctly."""
        inner = tvm_ffi.Dict({"x": 1})
        outer = tvm_ffi.Dict({"inner": inner})
        result = _roundtrip(outer)
        assert result["inner"]["x"] == 1

    def test_dict_with_array_value(self) -> None:
        """Dict with array values roundtrips correctly."""
        arr = tvm_ffi.convert([10, 20])
        d = tvm_ffi.Dict({"nums": arr})
        result = _roundtrip(d)
        assert list(result["nums"]) == [10, 20]


class TestShape:
    """Roundtrip tests for ffi.Shape containers."""

    def test_empty(self) -> None:
        """Empty shape roundtrips correctly."""
        shape = tvm_ffi.Shape(())
        _assert_roundtrip_eq(shape)

    def test_1d(self) -> None:
        """1D shape roundtrips correctly."""
        shape = tvm_ffi.Shape((10,))
        result = _roundtrip(shape)
        assert list(result) == [10]

    def test_nd(self) -> None:
        """N-D shape roundtrips correctly."""
        shape = tvm_ffi.Shape((1, 2, 3, 4))
        result = _roundtrip(shape)
        assert list(result) == [1, 2, 3, 4]


# ---------------------------------------------------------------------------
# Objects with reflection
# ---------------------------------------------------------------------------
class TestObjectSerialization:
    """Roundtrip tests for objects with reflection metadata."""

    def test_int_pair_roundtrip(self) -> None:
        """TestIntPair has refl::init and POD int64 fields."""
        pair = tvm_ffi.testing.TestIntPair(3, 7)
        s = to_json_graph_str(pair)
        result = from_json_graph_str(s)
        assert result.a == 3
        assert result.b == 7

    def test_int_pair_zero_values(self) -> None:
        """TestIntPair with zero values roundtrips correctly."""
        pair = tvm_ffi.testing.TestIntPair(0, 0)
        result = _roundtrip(pair)
        assert result.a == 0
        assert result.b == 0

    def test_int_pair_negative_values(self) -> None:
        """TestIntPair with negative values roundtrips correctly."""
        pair = tvm_ffi.testing.TestIntPair(-100, -200)
        result = _roundtrip(pair)
        assert result.a == -100
        assert result.b == -200

    def test_int_pair_large_values(self) -> None:
        """TestIntPair with large values roundtrips correctly."""
        pair = tvm_ffi.testing.TestIntPair(10**15, -(10**15))
        result = _roundtrip(pair)
        assert result.a == 10**15
        assert result.b == -(10**15)


# ---------------------------------------------------------------------------
# JSON structure verification
# ---------------------------------------------------------------------------
class TestJSONStructure:
    """Tests verifying the internal JSON graph structure."""

    def test_null_json_structure(self) -> None:
        """None produces a single node with type 'None'."""
        s = to_json_graph_str(None)
        parsed = json.loads(s)
        assert parsed["root_index"] == 0
        assert len(parsed["nodes"]) == 1
        assert parsed["nodes"][0]["type"] == "None"

    def test_bool_json_structure(self) -> None:
        """Bool produces a node with type 'bool'."""
        s = to_json_graph_str(True)
        parsed = json.loads(s)
        assert parsed["nodes"][0]["type"] == "bool"
        assert parsed["nodes"][0]["data"] is True

    def test_int_json_structure(self) -> None:
        """Int produces a node with type 'int'."""
        s = to_json_graph_str(42)
        parsed = json.loads(s)
        assert parsed["nodes"][0]["type"] == "int"
        assert parsed["nodes"][0]["data"] == 42

    def test_float_json_structure(self) -> None:
        """Float produces a node with type 'float'."""
        s = to_json_graph_str(3.14)
        parsed = json.loads(s)
        assert parsed["nodes"][0]["type"] == "float"
        assert parsed["nodes"][0]["data"] == pytest.approx(3.14)

    def test_string_json_structure(self) -> None:
        """String produces a node with type 'ffi.String'."""
        s = to_json_graph_str(tvm_ffi.convert("hello"))
        parsed = json.loads(s)
        assert parsed["nodes"][parsed["root_index"]]["type"] == "ffi.String"
        assert parsed["nodes"][parsed["root_index"]]["data"] == "hello"

    def test_array_json_structure(self) -> None:
        """Array data contains node index references."""
        s = to_json_graph_str(tvm_ffi.convert([1, 2]))
        parsed = json.loads(s)
        root = parsed["nodes"][parsed["root_index"]]
        assert root["type"] == "ffi.Array"
        # data should be list of node indices
        assert isinstance(root["data"], list)
        assert len(root["data"]) == 2

    def test_map_json_structure(self) -> None:
        """Map data contains flattened key-value node index pairs."""
        s = to_json_graph_str(tvm_ffi.convert({"a": 1}))
        parsed = json.loads(s)
        root = parsed["nodes"][parsed["root_index"]]
        assert root["type"] == "ffi.Map"
        assert isinstance(root["data"], list)
        # key-value pairs flattened: [key_idx, val_idx]
        assert len(root["data"]) == 2

    def test_dict_json_structure(self) -> None:
        """Dict data contains flattened key-value node index pairs."""
        s = to_json_graph_str(tvm_ffi.Dict({"a": 1}))
        parsed = json.loads(s)
        root = parsed["nodes"][parsed["root_index"]]
        assert root["type"] == "ffi.Dict"
        assert isinstance(root["data"], list)
        # key-value pairs flattened: [key_idx, val_idx]
        assert len(root["data"]) == 2

    def test_node_dedup(self) -> None:
        """Duplicate values should share the same node index."""
        s = to_json_graph_str(tvm_ffi.convert([42, 42, 42]))
        parsed = json.loads(s)
        root = parsed["nodes"][parsed["root_index"]]
        # all three elements should reference the same node
        assert root["data"][0] == root["data"][1] == root["data"][2]

    def test_object_pod_fields_are_inlined(self) -> None:
        """POD fields (int, bool, float) are inlined directly via field_static_type_index."""
        pair = tvm_ffi.testing.TestIntPair(3, 7)
        s = to_json_graph_str(pair)
        parsed = json.loads(s)
        root = parsed["nodes"][parsed["root_index"]]
        assert root["type"] == "testing.TestIntPair"
        # fields 'a' and 'b' should be inlined as direct int values (not node indices)
        assert root["data"]["a"] == 3
        assert root["data"]["b"] == 7


# ---------------------------------------------------------------------------
# Metadata
# ---------------------------------------------------------------------------
class TestMetadata:
    """Tests for optional metadata in serialized output."""

    def test_with_metadata(self) -> None:
        """Metadata dict appears in serialized JSON when provided."""
        s = to_json_graph_str(42, {"version": "1.0"})
        parsed = json.loads(s)
        assert "metadata" in parsed
        assert parsed["metadata"]["version"] == "1.0"

    def test_without_metadata(self) -> None:
        """Metadata key is absent when not provided."""
        s = to_json_graph_str(42)
        parsed = json.loads(s)
        assert "metadata" not in parsed


# ---------------------------------------------------------------------------
# Error cases
# ---------------------------------------------------------------------------
class TestErrors:
    """Tests for error handling in deserialization."""

    def test_invalid_json_string(self) -> None:
        """Invalid JSON string raises an error."""
        with pytest.raises(Exception):
            from_json_graph_str("not valid json")

    def test_empty_json_string(self) -> None:
        """Empty string raises an error."""
        with pytest.raises(Exception):
            from_json_graph_str("")

    def test_missing_root_index(self) -> None:
        """JSON without root_index raises an error."""
        with pytest.raises(Exception):
            from_json_graph_str('{"nodes": []}')

    def test_missing_nodes(self) -> None:
        """JSON without nodes raises an error."""
        with pytest.raises(Exception):
            from_json_graph_str('{"root_index": 0}')