"""Unit tests for mca_sdk.utils.serialization module."""

import json
import logging
import pytest
from dataclasses import dataclass
from unittest.mock import MagicMock, patch

from mca_sdk.utils.serialization import (
    serialize_for_telemetry,
    _json_default,
    extract_model_payload,
    _get_shape,
)


class TestSerializeForTelemetry:
    """Test suite for serialize_for_telemetry function."""

    def test_serialize_simple_dict(self):
        """Test that simple dicts are returned unchanged (OTel-compatible type)."""
        data = {"key": "value", "number": 42}
        result = serialize_for_telemetry(data)
        assert result == json.dumps(data)

    def test_serialize_simple_list(self):
        """Test that simple lists are returned unchanged (OTel-compatible type)."""
        data = [1, 2, 3, 4, 5]
        result = serialize_for_telemetry(data)
        assert result == json.dumps(data)

    def test_serialize_string(self):
        """Test that strings are returned unchanged (OTel-compatible type)."""
        data = "hello world"
        result = serialize_for_telemetry(data)
        assert result == "hello world"

    def test_serialize_number(self):
        """Test that numbers are returned unchanged (OTel-compatible type)."""
        assert serialize_for_telemetry(42) == 42
        assert serialize_for_telemetry(3.14) == 3.14

    def test_serialize_boolean(self):
        """Test that booleans are returned unchanged (OTel-compatible type)."""
        assert serialize_for_telemetry(True) is True
        assert serialize_for_telemetry(False) is False

    def test_serialize_none(self):
        """Test that None is returned unchanged (OTel-compatible type)."""
        assert serialize_for_telemetry(None) is None

    def test_serialize_nested_structure(self):
        """Test that nested dicts are returned unchanged (OTel-compatible type)."""
        data = {
            "users": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}],
            "metadata": {"count": 2, "active": True},
        }
        result = serialize_for_telemetry(data)
        assert result == json.dumps(data)

    def test_truncate_large_payload(self):
        """Test truncation of payload exceeding max_length."""
        data = "a" * 1000
        result = serialize_for_telemetry(data, max_length=100)
        assert len(result) == 100
        assert result.endswith("...[truncated]")

    def test_truncate_with_custom_marker(self):
        """Test truncation with custom truncation marker."""
        data = "x" * 500
        result = serialize_for_telemetry(data, max_length=50, truncate_marker="...")
        assert len(result) == 50
        assert result.endswith("...")

    def test_truncate_marker_longer_than_max_length(self):
        """Test edge case where truncation marker is longer than max_length."""
        data = "x" * 20  # Long enough to exceed max_length=5 and trigger truncation
        result = serialize_for_telemetry(data, max_length=5, truncate_marker="...[truncated]")
        assert result == "...[truncated]"

    def test_no_truncation_under_limit(self):
        """Test that data under max_length is not truncated (custom objects)."""

        class SmallObj:
            def __init__(self):
                self.key = "value"

        obj = SmallObj()
        result = serialize_for_telemetry(obj, max_length=1000)
        assert isinstance(result, str)
        assert "truncated" not in result

    def test_fallback_to_str_on_serialization_error(self):
        """Test fallback to str() when JSON serialization fails."""

        class NonSerializable:
            def __str__(self):
                return "NonSerializable instance"

            def __repr__(self):
                return "NonSerializable instance"

        obj = NonSerializable()
        result = serialize_for_telemetry(obj)
        # Should serialize using __dict__ which contains the __str__ method
        assert isinstance(result, str)

    def test_unicode_characters_preserved(self):
        """Test that Unicode characters are preserved in dicts returned unchanged."""
        data = {"message": "Hello 世界 🌍"}
        result = serialize_for_telemetry(data)
        assert json.loads(result)["message"] == "Hello 世界 🌍"

    def test_empty_dict(self):
        """Test that empty dict is returned unchanged."""
        assert serialize_for_telemetry({}) == "{}"

    def test_empty_list(self):
        """Test that empty list is returned unchanged."""
        assert serialize_for_telemetry([]) == "[]"


class TestJsonDefault:
    """Test suite for _json_default custom JSON encoder."""

    def test_numpy_array_with_tolist(self):
        """Test serialization of object with tolist() method (NumPy-like)."""

        class NumpyLike:
            def tolist(self):
                return [1, 2, 3, 4, 5]

        obj = NumpyLike()
        result = _json_default(obj)
        assert result == [1, 2, 3, 4, 5]

    def test_pandas_dataframe_with_to_dict(self):
        """Test serialization of object with to_dict() method (Pandas-like)."""

        class PandasLike:
            def to_dict(self):
                return {"col1": [1, 2, 3], "col2": [4, 5, 6]}

        obj = PandasLike()
        result = _json_default(obj)
        assert result == {"col1": [1, 2, 3], "col2": [4, 5, 6]}

    def test_dataclass_serialization(self):
        """Test serialization of dataclass objects."""

        @dataclass
        class Person:
            name: str
            age: int

        person = Person(name="Alice", age=30)
        result = _json_default(person)
        assert result == {"name": "Alice", "age": 30}

    def test_custom_class_with_dict(self):
        """Test serialization of custom class using __dict__."""

        class CustomClass:
            def __init__(self):
                self.attr1 = "value1"
                self.attr2 = 42

        obj = CustomClass()
        result = _json_default(obj)
        assert result == {"attr1": "value1", "attr2": 42}

    def test_fallback_to_str(self):
        """Test final fallback to str() when no serialization method works."""

        # Use a builtin type that doesn't have serialization methods
        class MinimalClass:
            pass

        obj = MinimalClass()
        result = _json_default(obj)
        # Should return the __dict__ or str representation
        assert isinstance(result, (dict, str))

    def test_tolist_failure_fallback(self, caplog):
        """Test fallback when tolist() raises exception."""

        class BrokenToList:
            def tolist(self):
                raise ValueError("tolist failed")

            def to_dict(self):
                return {"fallback": "to_dict"}

        obj = BrokenToList()
        with caplog.at_level(logging.WARNING):
            result = _json_default(obj)

        assert result == {"fallback": "to_dict"}
        assert "Failed to serialize" in caplog.text
        assert "tolist()" in caplog.text

    def test_to_dict_failure_fallback(self, caplog):
        """Test fallback when to_dict() raises exception."""

        class BrokenToDict:
            def to_dict(self):
                raise TypeError("to_dict failed")

        obj = BrokenToDict()
        with caplog.at_level(logging.WARNING):
            result = _json_default(obj)

        # Should fall back to __dict__
        assert isinstance(result, dict)
        assert "Failed to serialize" in caplog.text

    def test_dataclass_failure_fallback(self, caplog):
        """Test fallback when dataclass serialization fails."""

        class BrokenDataclass:
            __dataclass_fields__ = {"field": "metadata"}

            def __getattribute__(self, name):
                if name == "field":
                    raise AttributeError("field access failed")
                return super().__getattribute__(name)

            __dict__ = {"fallback": "dict"}

        obj = BrokenDataclass()
        with caplog.at_level(logging.WARNING):
            result = _json_default(obj)

        # Should fall back to __dict__
        assert isinstance(result, dict)

    def test_dict_failure_final_fallback(self, caplog):
        """Test str() fallback when no other method works."""

        # Create an object without common serialization methods
        obj = object()  # Builtin object type

        with caplog.at_level(logging.DEBUG):
            result = _json_default(obj)

        # Should use str() fallback for object()
        assert isinstance(result, str)
        assert "object" in result

    def test_logging_debug_messages(self, caplog):
        """Test that debug logging works for successful serializations."""

        class SimpleClass:
            def tolist(self):
                return [1, 2, 3]

        obj = SimpleClass()
        with caplog.at_level(logging.DEBUG):
            _json_default(obj)

        assert "Serialized" in caplog.text
        assert "tolist()" in caplog.text

    def test_mixed_serialization_methods(self):
        """Test object with multiple serialization methods uses first available."""

        class MultiMethod:
            def tolist(self):
                return [1, 2, 3]

            def to_dict(self):
                return {"should": "not reach this"}

        obj = MultiMethod()
        result = _json_default(obj)
        # Should use tolist() since it's checked first
        assert result == [1, 2, 3]


class TestIntegrationScenarios:
    """Integration tests for real-world serialization scenarios."""

    def test_complex_nested_with_custom_objects(self):
        """Test serialization of custom objects (non-dict) with nested attrs."""

        class Model:
            def __init__(self, name, version):
                self.name = name
                self.version = version

        # Pass the custom object directly (not wrapped in a dict) to trigger _json_default
        obj = Model("classifier", "1.0")
        result = serialize_for_telemetry(obj)

        parsed = json.loads(result)
        assert parsed["name"] == "classifier"
        assert parsed["version"] == "1.0"

    def test_large_dataset_truncation_preserves_structure(self):
        """Test that truncation happens at correct position for custom objects."""

        class LargeObj:
            def __init__(self):
                self.data = "x" * 10000

        obj = LargeObj()
        result = serialize_for_telemetry(obj, max_length=100)

        assert len(result) == 100
        assert result.startswith('{"data"')
        assert result.endswith("...[truncated]")

    def test_serialization_exception_types(self):
        """Test handling of different JSON serialization exception types."""

        # Object that has __dict__ will be serialized successfully
        class RegularObj:
            def __init__(self):
                self.value = "handled"

        obj = RegularObj()
        result = serialize_for_telemetry(obj)

        # Should serialize successfully using __dict__
        assert "value" in result
        assert "handled" in result


class TestGetShape:
    """Test suite for _get_shape helper function."""

    def test_get_shape_list(self):
        """Test shape extraction from list."""
        data = [1, 2, 3, 4, 5]
        result = _get_shape(data)
        assert result == "(5,)"

    def test_get_shape_tuple(self):
        """Test shape extraction from tuple."""
        data = (1, 2, 3)
        result = _get_shape(data)
        assert result == "(3,)"

    def test_get_shape_with_shape_attribute(self):
        """Test shape extraction from object with shape attribute (NumPy-like)."""

        class NumpyLike:
            shape = (10, 5)

        obj = NumpyLike()
        result = _get_shape(obj)
        assert result == "(10, 5)"

    def test_get_shape_string_returns_none(self):
        """Test that strings don't have shape."""
        result = _get_shape("hello")
        assert result is None

    def test_get_shape_dict_returns_none(self):
        """Test that dicts don't have shape."""
        result = _get_shape({"key": "value"})
        assert result is None

    def test_get_shape_exception_returns_none(self):
        """Test that exceptions during shape extraction return None."""

        class BrokenShape:
            @property
            def shape(self):
                raise ValueError("Broken shape")

        obj = BrokenShape()
        result = _get_shape(obj)
        assert result is None


class TestExtractModelPayload:
    """Test suite for extract_model_payload function."""

    def test_extract_simple_list(self):
        """Test extraction of simple list with shape."""
        data = [1, 2, 3, 4, 5]
        result = extract_model_payload(data)

        assert result["data"] == "[1, 2, 3, 4, 5]"
        assert result["shape"] == "(5,)"
        assert result["truncated"] is False

    def test_extract_with_numpy_like_shape(self):
        """Test extraction with NumPy-like array."""

        class NumpyLike:
            shape = (100, 5)

            def __str__(self):
                return "[[0 1 2 3 4]\n [5 6 7 8 9]\n ...]"

        data = NumpyLike()
        result = extract_model_payload(data)

        assert "[[0 1 2 3 4]" in result["data"]
        assert result["shape"] == "(100, 5)"
        assert result["truncated"] is False

    def test_extract_pandas_like_dataframe(self):
        """Test extraction with Pandas-like DataFrame."""

        class DataFrameLike:
            shape = (10, 3)

            def __str__(self):
                return "   col1  col2  col3\n0     1     2     3"

        data = DataFrameLike()
        result = extract_model_payload(data)

        assert "col1" in result["data"]
        assert result["shape"] == "(10, 3)"
        assert result["truncated"] is False

    def test_extract_string_no_shape(self):
        """Test extraction of string (no shape)."""
        data = "simple string"
        result = extract_model_payload(data)

        assert result["data"] == "simple string"
        assert result["shape"] is None
        assert result["truncated"] is False

    def test_extract_dict_no_shape(self):
        """Test extraction of dict (no shape)."""
        data = {"key": "value", "number": 42}
        result = extract_model_payload(data)

        assert "key" in result["data"]
        assert "value" in result["data"]
        assert result["shape"] is None
        assert result["truncated"] is False

    def test_truncation_by_bytes(self):
        """Test byte-level truncation for large data."""
        # Use smaller multiplier to avoid object size check triggering
        data = "x" * 5000
        result = extract_model_payload(data, max_size_bytes=1000, max_object_size_multiplier=100)

        assert len(result["data"]) < 1100  # Should be around 1000 + truncation message
        assert "truncated" in result["data"]
        assert " bytes total]" in result["data"]
        assert result["truncated"] is True

    def test_oversized_object_prevention(self):
        """Test OOM prevention for huge objects."""
        # sys.getsizeof() reports the size of the string object itself (including data)
        # Create a string that's larger than max_size_bytes * max_object_size_multiplier
        huge_string = "x" * 150000  # 150KB string
        result = extract_model_payload(
            huge_string, max_size_bytes=1000, max_object_size_multiplier=10
        )

        # Should trigger object size check (150000 > 1000 * 10) and skip serialization
        assert "[Object too large to serialize:" in result["data"]
        assert result["truncated"] is True

    def test_utf8_handling(self):
        """Test proper UTF-8 handling during truncation."""
        # Create string with multi-byte UTF-8 characters
        data = "Hello 世界 " * 100  # Each 世界 is 3 bytes in UTF-8
        result = extract_model_payload(data, max_size_bytes=50)

        # Should truncate without breaking UTF-8 encoding
        assert result["truncated"] is True
        # Result should be valid UTF-8 (no decode errors)
        assert isinstance(result["data"], str)

    def test_returns_dict_structure(self):
        """Test that function returns dict with expected keys."""
        data = [1, 2, 3]
        result = extract_model_payload(data)

        assert isinstance(result, dict)
        assert "data" in result
        assert "shape" in result
        assert "truncated" in result

    def test_exception_handling(self):
        """Test graceful handling of serialization exceptions."""

        class ProblematicObject:
            def __str__(self):
                raise RuntimeError("Cannot convert to string")

        data = ProblematicObject()
        result = extract_model_payload(data)

        assert result["data"] == "[Failed to serialize data]"
        assert result["truncated"] is True

    def test_custom_max_size(self):
        """Test custom max_size_bytes parameter."""
        data = "a" * 500
        result = extract_model_payload(data, max_size_bytes=100)

        assert result["truncated"] is True
        assert len(result["data"]) < 150  # Around 100 + truncation message

    def test_empty_data(self):
        """Test handling of empty data structures."""
        # Empty list
        result = extract_model_payload([])
        assert result["data"] == "[]"
        assert result["shape"] == "(0,)"
        assert result["truncated"] is False

        # Empty string
        result = extract_model_payload("")
        assert result["data"] == ""
        assert result["shape"] is None
        assert result["truncated"] is False

    def test_none_data(self):
        """Test handling of None data."""
        result = extract_model_payload(None)
        assert result["data"] == "None"
        assert result["shape"] is None
        assert result["truncated"] is False

    def test_nested_list_shape(self):
        """Test Issue #4 Fix: Deep shape extraction for nested lists."""
        data = [[1, 2, 3], [4, 5, 6]]
        result = extract_model_payload(data)
        assert result["shape"] == "(2, 3)", f"Expected (2, 3), got {result['shape']}"
        assert result["truncated"] is False

    def test_multi_arg_dict_unwrapping(self):
        """Test Issue #5 Fix: Unwrap {"args": [...], "kwargs": {...}} for shape."""
        data = {"args": ([1, 2, 3, 4, 5],), "kwargs": {}}
        result = extract_model_payload(data)
        # Should extract shape from args[0]
        assert result["shape"] == "(5,)", f"Expected (5,), got {result['shape']}"

    def test_kwargs_only_dict_unwrapping(self):
        """Test Issue #5 Complete Fix: Extract shape from kwargs when args is empty."""
        # Common pattern: model.predict(X=data)
        data = {"args": (), "kwargs": {"X": [[1, 2, 3], [4, 5, 6]]}}
        result = extract_model_payload(data)
        # Should extract shape from kwargs["X"]
        assert result["shape"] == "(2, 3)", f"Expected (2, 3), got {result['shape']}"

    def test_kwargs_only_from_autolog(self):
        """Test Issue #8 Fix: Handle {"kwargs": {...}} structure from autolog."""
        # When autolog captures only kwargs
        data = {"kwargs": {"X": [1, 2, 3, 4, 5]}}
        result = extract_model_payload(data)
        # Should extract shape from kwargs["X"]
        assert result["shape"] == "(5,)", f"Expected (5,), got {result['shape']}"

    def test_kwargs_with_non_standard_param_name(self):
        """Test kwargs extraction with non-standard parameter name."""
        # If parameter isn't in common list, use first kwarg
        data = {"args": (), "kwargs": {"my_custom_input": [[1, 2], [3, 4]]}}
        result = extract_model_payload(data)
        # Should still extract shape from first kwarg
        assert result["shape"] == "(2, 2)", f"Expected (2, 2), got {result['shape']}"

    def test_truncation_respects_byte_limit(self):
        """Test Issue #7 Fix: Truncation message included in max_size_bytes."""
        data = "x" * 5000
        max_size = 100
        result = extract_model_payload(data, max_size_bytes=max_size)

        # Total result should not exceed max_size_bytes
        result_bytes = result["data"].encode("utf-8")
        assert (
            len(result_bytes) <= max_size
        ), f"Result {len(result_bytes)} bytes exceeds limit of {max_size} bytes"
        assert result["truncated"] is True

    def test_generator_handling(self):
        """Test Issue #9 Fix: Detect generators without consuming them."""

        def sample_generator():
            yield 1
            yield 2
            yield 3

        gen = sample_generator()
        result = extract_model_payload(gen)

        assert result["data"] == "[Unconsumable Iterator/Generator]"
        assert result["truncated"] is True
        # Generator should not be consumed
        assert next(gen) == 1  # Still works

    def test_large_collection_prevention(self):
        """Test Issue #3 Fix: Prevent serialization of huge collections."""
        # Create a list larger than MAX_COLLECTION_LENGTH
        data = [0] * 150000
        result = extract_model_payload(data, max_size_bytes=1000)

        assert "[Collection too large to serialize:" in result["data"]
        assert result["truncated"] is True

    def test_serialization_error_logging_level(self):
        """Test Issue #11 Fix: Errors logged at warning level."""
        import logging

        # Mock logger to capture warning calls
        with MagicMock() as mock_logger:
            import mca_sdk.utils.serialization as ser_module

            original_logger = ser_module.logger
            ser_module.logger = mock_logger

            class ProblematicObject:
                def __str__(self):
                    raise RuntimeError("Cannot convert")

            data = ProblematicObject()
            result = extract_model_payload(data)

            # Should log at warning level, not debug
            mock_logger.warning.assert_called()
            ser_module.logger = original_logger

        assert result["truncated"] is True
