tests/test_transforms/test_universal_json_crush.py

"""Tests for universal JSON compression (all array types).

Verifies that SmartCrusher handles arrays of dicts, strings, numbers,
mixed types, and nested arrays — with consistent safety guarantees
across all types.
"""

from __future__ import annotations

import json

import pytest

from headroom.transforms.smart_crusher import (
    ArrayType,
    SmartCrusher,
    SmartCrusherConfig,
    _classify_array,
)

# =====================================================================
# Fixtures
# =====================================================================


@pytest.fixture
def crusher():
    """SmartCrusher configured for testing."""
    return SmartCrusher(
        config=SmartCrusherConfig(
            min_items_to_analyze=5,
            min_tokens_to_crush=0,  # Always crush
            max_items_after_crush=15,
        )
    )


@pytest.fixture
def crusher_large_k():
    """SmartCrusher with higher max items for larger test arrays."""
    return SmartCrusher(
        config=SmartCrusherConfig(
            min_items_to_analyze=5,
            min_tokens_to_crush=0,
            max_items_after_crush=50,
        )
    )


# =====================================================================
# Type Classification
# =====================================================================


class TestClassifyArray:
    def test_dict_array(self):
        assert _classify_array([{"a": 1}, {"b": 2}]) == ArrayType.DICT_ARRAY

    def test_string_array(self):
        assert _classify_array(["hello", "world", "foo"]) == ArrayType.STRING_ARRAY

    def test_number_array_int(self):
        assert _classify_array([1, 2, 3]) == ArrayType.NUMBER_ARRAY

    def test_number_array_float(self):
        assert _classify_array([1.0, 2.5, 3.7]) == ArrayType.NUMBER_ARRAY

    def test_number_array_mixed_int_float(self):
        assert _classify_array([1, 2.5, 3]) == ArrayType.NUMBER_ARRAY

    def test_bool_array(self):
        assert _classify_array([True, False, True]) == ArrayType.BOOL_ARRAY

    def test_nested_array(self):
        assert _classify_array([[1, 2], [3, 4]]) == ArrayType.NESTED_ARRAY

    def test_mixed_array(self):
        assert _classify_array([{"a": 1}, "string", 42]) == ArrayType.MIXED_ARRAY

    def test_empty(self):
        assert _classify_array([]) == ArrayType.EMPTY

    def test_single_dict(self):
        assert _classify_array([{"key": "val"}]) == ArrayType.DICT_ARRAY

    def test_single_string(self):
        assert _classify_array(["only"]) == ArrayType.STRING_ARRAY

    def test_none_values(self):
        # Arrays with None mixed in are MIXED
        assert _classify_array([1, None, 3]) == ArrayType.MIXED_ARRAY

    def test_bool_not_confused_with_int(self):
        # Python's True/False are int subclasses — we handle this
        assert _classify_array([True, False]) == ArrayType.BOOL_ARRAY
        # But mixed bools and ints should be MIXED or NUMBER depending on impl
        result = _classify_array([True, 42])
        assert result in (ArrayType.MIXED_ARRAY, ArrayType.NUMBER_ARRAY)


# =====================================================================
# String Array Compression
# =====================================================================


class TestCrushStringArray:
    def test_basic_compression(self, crusher):
        strings = [f"item_{i}" for i in range(100)]
        crushed, strategy = crusher._crush_string_array(strings)
        assert len(crushed) < len(strings)
        assert "string:adaptive" in strategy

    def test_errors_always_preserved(self, crusher):
        strings = ["ok"] * 50 + ["error: connection timeout", "failed: auth denied"] + ["ok"] * 48
        crushed, strategy = crusher._crush_string_array(strings)
        assert any("error" in s for s in crushed)
        assert any("failed" in s for s in crushed)

    def test_first_last_kept(self, crusher):
        strings = [f"item_{i}" for i in range(50)]
        crushed, strategy = crusher._crush_string_array(strings)
        # First item always present
        assert strings[0] in crushed
        # Last item always present
        assert strings[-1] in crushed

    def test_dedup_reduces_output(self, crusher):
        # 95 identical + 5 unique
        strings = ["repeated_value"] * 95 + [f"unique_{i}" for i in range(5)]
        crushed, strategy = crusher._crush_string_array(strings)
        # Should massively reduce — not keep 95 copies
        assert len(crushed) < 20
        # All 5 unique values should survive (they have high info value)
        for i in range(5):
            assert f"unique_{i}" in crushed

    def test_below_threshold_passthrough(self, crusher):
        strings = ["a", "b", "c"]  # Below min_items_to_analyze=5
        # Direct method call — should passthrough since <= 8
        crushed, strategy = crusher._crush_string_array(strings)
        assert crushed == strings
        assert "passthrough" in strategy

    def test_empty_strings_handled(self, crusher):
        strings = [""] * 20
        crushed, strategy = crusher._crush_string_array(strings)
        # Should not crash
        assert isinstance(crushed, list)

    def test_unicode_strings(self, crusher):
        strings = [f"日本語テスト_{i}" for i in range(50)]
        crushed, strategy = crusher._crush_string_array(strings)
        assert len(crushed) < len(strings)
        assert all(isinstance(s, str) for s in crushed)

    def test_length_anomalies_preserved(self, crusher_large_k):
        # Most strings are short, one is very long
        strings = ["short"] * 95 + ["x" * 10000] + ["short"] * 4
        crushed, strategy = crusher_large_k._crush_string_array(strings)
        assert any(len(s) > 1000 for s in crushed)


# =====================================================================
# Number Array Compression
# =====================================================================


class TestCrushNumberArray:
    def test_basic_compression(self, crusher):
        numbers = [42.0 + i * 0.1 for i in range(100)]
        crushed, strategy = crusher._crush_number_array(numbers)
        assert len(crushed) < len(numbers)
        assert "number:adaptive" in strategy

    def test_summary_prepended(self, crusher):
        numbers = list(range(100))
        crushed, strategy = crusher._crush_number_array(numbers)
        # First element should be the stats summary string
        assert isinstance(crushed[0], str)
        assert "numbers:" in crushed[0]
        assert "min=" in crushed[0]
        assert "max=" in crushed[0]

    def test_outliers_preserved(self, crusher):
        # Normal values around 50 with one extreme outlier
        numbers = [50.0 + i * 0.01 for i in range(100)] + [999.9]
        crushed, strategy = crusher._crush_number_array(numbers)
        assert 999.9 in crushed
        assert "outliers" in strategy

    def test_all_identical(self, crusher):
        numbers = [42.0] * 100
        crushed, strategy = crusher._crush_number_array(numbers)
        # With all identical, should compress heavily
        # Summary + a few representatives
        numeric_values = [v for v in crushed if isinstance(v, (int, float))]
        assert all(v == 42.0 for v in numeric_values)

    def test_first_last_kept(self, crusher):
        numbers = list(range(50))
        crushed, strategy = crusher._crush_number_array(numbers)
        numeric_values = [v for v in crushed if isinstance(v, (int, float))]
        assert 0 in numeric_values  # First
        assert 49 in numeric_values  # Last

    def test_change_point_preserved(self, crusher_large_k):
        # Stable at 10, then jumps to 100
        numbers = [10.0] * 50 + [100.0] * 50
        crushed, strategy = crusher_large_k._crush_number_array(numbers)
        numeric_values = [v for v in crushed if isinstance(v, (int, float))]
        # Both 10.0 and 100.0 should be present
        assert 10.0 in numeric_values
        assert 100.0 in numeric_values

    def test_nan_inf_filtered(self, crusher):
        numbers = [1.0, 2.0, float("nan"), float("inf"), 3.0] * 10
        crushed, strategy = crusher._crush_number_array(numbers)
        # Should not crash; stats should be based on finite values
        assert isinstance(crushed[0], str)

    def test_integers_preserved_as_int(self, crusher):
        numbers = list(range(50))
        crushed, strategy = crusher._crush_number_array(numbers)
        numeric_values = [v for v in crushed if isinstance(v, (int, float))]
        # Integers should remain integers (not converted to float)
        assert any(isinstance(v, int) for v in numeric_values)

    def test_statistics_accuracy(self, crusher):
        numbers = list(range(1, 101))  # 1 to 100
        crushed, strategy = crusher._crush_number_array(numbers)
        summary = crushed[0]
        assert "min=1" in summary
        assert "max=100" in summary
        assert "mean=50.5" in summary


# =====================================================================
# Mixed Array Compression
# =====================================================================


class TestCrushMixedArray:
    def test_basic_compression(self, crusher_large_k):
        mixed = [{"id": i} for i in range(30)] + [f"msg_{i}" for i in range(30)]
        crushed, strategy = crusher_large_k._crush_mixed_array(mixed)
        # With diversity-aware K, unique items may all be kept.
        # Verify compression happened OR items are preserved due to high diversity.
        assert len(crushed) <= len(mixed)
        assert "mixed" in strategy

    def test_small_groups_kept(self, crusher):
        # 50 dicts + 3 strings (below threshold)
        mixed = [{"id": i} for i in range(50)] + ["rare_1", "rare_2", "rare_3"]
        crushed, strategy = crusher._crush_mixed_array(mixed)
        # All 3 rare strings should be kept (below min_items threshold)
        assert "rare_1" in crushed
        assert "rare_2" in crushed
        assert "rare_3" in crushed

    def test_errors_across_types(self, crusher_large_k):
        mixed = (
            [{"status": "ok"}] * 30
            + [{"status": "error: timeout"}]
            + ["error: auth failed"]
            + [f"ok_{i}" for i in range(30)]
        )
        crushed, strategy = crusher_large_k._crush_mixed_array(mixed)
        crushed_str = json.dumps(crushed)
        assert "error: timeout" in crushed_str
        assert "error: auth failed" in crushed_str

    def test_original_order_preserved(self, crusher_large_k):
        mixed = [{"id": i} for i in range(20)] + [f"str_{i}" for i in range(20)]
        crushed, strategy = crusher_large_k._crush_mixed_array(mixed)
        # Dicts should come before strings (original order)
        first_str_idx = next(
            (i for i, item in enumerate(crushed) if isinstance(item, str)), len(crushed)
        )
        last_dict_idx = max(
            (i for i, item in enumerate(crushed) if isinstance(item, dict)), default=-1
        )
        assert last_dict_idx < first_str_idx

    def test_passthrough_small(self, crusher):
        mixed = [1, "two", {"three": 3}]
        crushed, strategy = crusher._crush_mixed_array(mixed)
        assert crushed == mixed
        assert "passthrough" in strategy


# =====================================================================
# Adaptive K
# =====================================================================


class TestAdaptiveK:
    def test_scales_with_n(self, crusher):
        """K grows sublinearly with collection size (or saturates at max_items)."""
        small = [f"item_{i}" for i in range(20)]
        large = [f"item_{i}" for i in range(500)]

        k_small = crusher._compute_k_split(small)[0]
        k_large = crusher._compute_k_split(large)[0]

        # With diversity-aware K, both may saturate at max_items_after_crush
        # for highly unique items. The key property: k never exceeds max_items.
        assert k_large >= k_small
        assert k_large <= crusher.config.max_items_after_crush

    def test_respects_max_items(self, crusher):
        items = [f"item_{i}" for i in range(1000)]
        k_total, _, _, _ = crusher._compute_k_split(items)
        assert k_total <= crusher.config.max_items_after_crush

    def test_first_last_fractions(self, crusher):
        items = [f"item_{i}" for i in range(100)]
        k_total, k_first, k_last, k_importance = crusher._compute_k_split(items)
        # First and last should be roughly the configured fractions
        assert k_first >= 1
        assert k_last >= 1
        assert k_first + k_last + k_importance == k_total

    def test_homogeneous_vs_diverse(self, crusher_large_k):
        """Homogeneous data should produce smaller K than diverse data."""
        homogeneous = ["same value"] * 100
        diverse = [f"unique_value_{i}_{'x' * (i * 10)}" for i in range(100)]

        k_homo = crusher_large_k._compute_k_split(homogeneous)[0]
        k_diverse = crusher_large_k._compute_k_split(diverse)[0]

        # Diverse should need more items (or at least equal)
        assert k_diverse >= k_homo


# =====================================================================
# Safety Guarantees (parametrized across types)
# =====================================================================


class TestSafetyGuarantees:
    @pytest.mark.parametrize(
        "items,item_type",
        [
            ([{"status": "ok"}] * 50 + [{"status": "error: timeout"}], "dict"),
            (["ok"] * 50 + ["error: connection failed"], "string"),
            (
                [{"id": i} for i in range(30)]
                + ["error: auth denied"]
                + [f"ok_{i}" for i in range(30)],
                "mixed",
            ),
        ],
        ids=["dict_array", "string_array", "mixed_array"],
    )
    def test_errors_never_dropped(self, crusher_large_k, items, item_type):
        """Error items must be preserved regardless of array type."""
        if item_type == "dict":
            crushed, _ = crusher_large_k._crush_string_array([json.dumps(i) for i in items])
            crushed_text = " ".join(crushed)
        elif item_type == "string":
            crushed, _ = crusher_large_k._crush_string_array(items)
            crushed_text = " ".join(crushed)
        elif item_type == "mixed":
            crushed, _ = crusher_large_k._crush_mixed_array(items)
            crushed_text = json.dumps(crushed)

        assert "error" in crushed_text.lower()

    @pytest.mark.parametrize(
        "items",
        [
            [f"item_{i}" for i in range(50)],
            list(range(50)),
        ],
        ids=["string_array", "number_array"],
    )
    def test_first_last_always_present(self, crusher, items):
        """First and last items must be present in output."""
        if isinstance(items[0], str):
            crushed, _ = crusher._crush_string_array(items)
            assert items[0] in crushed
            assert items[-1] in crushed
        else:
            crushed, _ = crusher._crush_number_array(items)
            numeric = [v for v in crushed if isinstance(v, (int, float))]
            assert items[0] in numeric
            assert items[-1] in numeric

    @pytest.mark.parametrize(
        "items",
        [
            ["a", "b", "c"],
            [1, 2, 3],
            [{"x": 1}, "y", 2],
        ],
        ids=["string_below", "number_below", "mixed_below"],
    )
    def test_passthrough_below_min_items(self, crusher, items):
        """Arrays below min_items_to_analyze pass through unchanged."""
        if all(isinstance(i, str) for i in items):
            crushed, strategy = crusher._crush_string_array(items)
        elif all(isinstance(i, (int, float)) for i in items):
            crushed, strategy = crusher._crush_number_array(items)
        else:
            crushed, strategy = crusher._crush_mixed_array(items)
        assert "passthrough" in strategy


# =====================================================================
# Integration: Full pipeline
# =====================================================================


class TestFullPipelineIntegration:
    """Test that new types work through the compress() API."""

    def test_string_array_via_compress(self):
        from headroom import compress

        strings = [f"log line {i}: GET /api 200" for i in range(100)]
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Show logs"},
            {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {"name": "get_logs", "arguments": "{}"},
                    }
                ],
            },
            {"role": "tool", "tool_call_id": "call_1", "content": json.dumps(strings)},
        ]
        result = compress(messages)
        assert result.tokens_saved > 0
        assert result.compression_ratio > 0

    def test_number_array_via_compress(self):
        from headroom import compress

        numbers = [42.0 + i * 0.1 for i in range(200)]
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Show metrics"},
            {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {"name": "get_metrics", "arguments": "{}"},
                    }
                ],
            },
            {"role": "tool", "tool_call_id": "call_1", "content": json.dumps(numbers)},
        ]
        result = compress(messages)
        assert result.tokens_saved > 0

    def test_dict_array_unchanged(self):
        """Verify dict arrays still work (regression test)."""
        from headroom import compress

        data = [{"id": i, "name": f"user_{i}", "status": "active"} for i in range(100)]
        messages = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "List users"},
            {
                "role": "assistant",
                "content": None,
                "tool_calls": [
                    {
                        "id": "call_1",
                        "type": "function",
                        "function": {"name": "list_users", "arguments": "{}"},
                    }
                ],
            },
            {"role": "tool", "tool_call_id": "call_1", "content": json.dumps(data)},
        ]
        result = compress(messages)
        assert result.tokens_saved > 0
        assert result.compression_ratio > 0