"""Additional branch-heavy tests for serialize helpers."""

from __future__ import annotations

import types
from typing import Any

import numpy as np
import pytest

from llmdebug.config import SnapshotConfig
from llmdebug.serialize import (
    _build_local_metadata_entry,
    _compute_array_stats_and_anomalies,
    _jax_stats_and_anomalies,
    _numpy_like_stats_and_anomalies,
    _sample_array_head,
    _torch_stats_and_anomalies,
    get_array_element_count,
    get_requires_grad,
    get_tensor_device,
    serialize_locals_bundle,
    summarize_dataframe,
    summarize_series,
    to_jsonlike,
)


def test_get_tensor_device_paths() -> None:
    class _CpuDevice:
        type = "cpu"

        def __str__(self) -> str:
            return "cpu"

    class _CudaDevice:
        type = "cuda"

        def __str__(self) -> str:
            return "cuda:0"

    class _WithCpuDevice:
        device = _CpuDevice()

    class _WithCudaDevice:
        device = _CudaDevice()

    class _WithDevices:
        def devices(self) -> set[str]:
            return {"gpu:0"}

    class _WithBadDevice:
        @property
        def device(self) -> object:
            raise RuntimeError("boom")

    assert get_tensor_device(_WithCpuDevice()) is None
    assert get_tensor_device(_WithCudaDevice()) == "cuda:0"
    assert get_tensor_device(_WithDevices()) == "{'gpu:0'}"
    assert get_tensor_device(_WithBadDevice()) is None


def test_get_requires_grad_paths() -> None:
    class _TrueGrad:
        requires_grad = True

    class _FalseGrad:
        requires_grad = False

    class _BadGrad:
        @property
        def requires_grad(self) -> bool:
            raise RuntimeError("boom")

    assert get_requires_grad(_TrueGrad()) is True
    assert get_requires_grad(_FalseGrad()) is None
    assert get_requires_grad(_BadGrad()) is None


def test_get_array_element_count_fallbacks() -> None:
    class _FromShape:
        size = "not-an-int"
        shape = (3, 4)

    class _NegativeShape:
        shape = (3, -1)

    class _UnknownShape:
        shape = ("?", 4)

    assert get_array_element_count(_FromShape()) == 12
    assert get_array_element_count(_NegativeShape()) is None
    assert get_array_element_count(_UnknownShape()) is None
    assert get_array_element_count(object()) is None


def test_sample_array_head_fallback_and_error_paths() -> None:
    class _RavelOnly:
        def reshape(self, _shape: int) -> np.ndarray:
            raise RuntimeError("no reshape")

        def ravel(self) -> np.ndarray:
            return np.array([1, 2, 3], dtype=np.int32)

    class _FlattenOnly:
        def flatten(self) -> list[int]:
            return [4, 5, 6]

    assert _sample_array_head(_RavelOnly(), max_items=2) == [1.0, 2.0]
    assert _sample_array_head(_FlattenOnly(), max_items=2) == [4.0, 5.0]

    with pytest.raises(TypeError, match="array does not support"):
        _sample_array_head(object(), max_items=2)


def test_compute_array_stats_and_anomalies_short_circuit_paths() -> None:
    arr = np.arange(10, dtype=np.float32)
    assert _compute_array_stats_and_anomalies(arr, max_elements=3) == (None, None)
    assert _compute_array_stats_and_anomalies(np.array([], dtype=np.float32)) == (None, None)

    class _Unknown:
        __module__ = "custom.module"

    assert _compute_array_stats_and_anomalies(_Unknown()) == (None, None)


def test_torch_stats_and_anomalies_with_stub(monkeypatch: pytest.MonkeyPatch) -> None:
    class _Scalar:
        def __init__(self, value: float) -> None:
            self._value = value

        def item(self) -> float:
            return self._value

    class _Tensor:
        def __init__(self, values: list[float]) -> None:
            self.values = np.asarray(values, dtype=np.float32)

        def detach(self) -> _Tensor:
            return self

        def cpu(self) -> _Tensor:
            return self

        def float(self) -> _Tensor:
            return self

    class _TensorMask:
        def __init__(self, mask: np.ndarray) -> None:
            self.mask = mask

        def sum(self) -> _Scalar:
            return _Scalar(float(np.sum(self.mask)))

    fake_torch = types.SimpleNamespace(
        isnan=lambda t: _TensorMask(np.isnan(t.values)),
        isinf=lambda t: _TensorMask(np.isinf(t.values)),
        min=lambda t: _Scalar(float(np.min(t.values))),
        max=lambda t: _Scalar(float(np.max(t.values))),
        mean=lambda t: _Scalar(float(np.mean(t.values))),
        std=lambda t, unbiased=False: _Scalar(float(np.std(t.values))),
    )
    monkeypatch.setitem(__import__("sys").modules, "torch", fake_torch)

    stats, anomalies = _torch_stats_and_anomalies(
        _Tensor([1.0, float("nan"), float("inf")]), is_numeric_dtype=True
    )
    assert stats is not None
    assert np.isnan(stats["min"])
    assert anomalies == {"nan": 1, "inf": 1}

    stats2, anomalies2 = _torch_stats_and_anomalies(_Tensor([1.0]), is_numeric_dtype=False)
    assert stats2 is None
    assert anomalies2 is None


def test_jax_stats_and_anomalies_with_stub(monkeypatch: pytest.MonkeyPatch) -> None:
    fake_jnp = types.SimpleNamespace(
        float32=np.float32,
        isnan=np.isnan,
        isinf=np.isinf,
        min=np.min,
        max=np.max,
        mean=np.mean,
        std=np.std,
    )
    fake_jax = types.SimpleNamespace(numpy=fake_jnp)
    monkeypatch.setitem(__import__("sys").modules, "jax", fake_jax)
    monkeypatch.setitem(__import__("sys").modules, "jax.numpy", fake_jnp)

    class _JaxArray:
        def __init__(self, values: list[float]) -> None:
            self._values = np.asarray(values, dtype=np.float32)

        def astype(self, dtype: Any) -> np.ndarray:
            return self._values.astype(dtype)

    stats, anomalies = _jax_stats_and_anomalies(
        _JaxArray([1.0, float("nan"), float("inf")]), is_numeric_dtype=True
    )
    assert stats is not None
    assert anomalies == {"nan": 1, "inf": 1}

    stats2, anomalies2 = _jax_stats_and_anomalies(_JaxArray([1.0]), is_numeric_dtype=False)
    assert stats2 is None
    assert anomalies2 is None


def test_numpy_like_stats_and_anomalies_cupy_path() -> None:
    class _FakeCupy:
        __module__ = "cupy.core.core"

        def __init__(self, values: list[float]) -> None:
            self._values = np.asarray(values, dtype=np.float32)
            self.dtype = self._values.dtype

        def get(self) -> np.ndarray:
            return self._values

    stats, anomalies = _numpy_like_stats_and_anomalies(
        _FakeCupy([1.0, 2.0, 3.0]),
        module="cupy.core.core",
        is_numeric_dtype=True,
    )
    assert stats is not None
    assert anomalies is None


def test_dataframe_and_series_summaries_with_fakes() -> None:
    class _Values:
        def tolist(self) -> list[list[int]]:
            return [[1, 2], [3, 4]]

    class _Head:
        values = _Values()

    class DataFrame:
        __module__ = "pandas.core.frame"

        def __init__(self) -> None:
            self.shape = (2, 2)
            self.columns = ["a", "b", "c"]

        def __len__(self) -> int:
            return 2

        def head(self, _rows: int) -> _Head:
            return _Head()

    class Series:
        __module__ = "pandas.core.series"

        def __init__(self) -> None:
            self.shape = (6,)
            self.name = "s"
            self.dtype = "int64"

        def head(self, _rows: int) -> Series:
            return self

        def tolist(self) -> list[int]:
            return [1, 2, 3, 4, 5]

        def __len__(self) -> int:
            return 6

    cfg = SnapshotConfig(max_items=2)
    df_summary = summarize_dataframe(DataFrame(), cfg)
    assert df_summary["__dataframe__"] is True
    assert df_summary["columns"] == ["a", "b"]
    assert df_summary["columns_truncated"] is True
    assert "head" in df_summary

    series_summary = summarize_series(Series(), cfg)
    assert series_summary["__series__"] is True
    assert series_summary["name"] == "s"
    assert series_summary["head_truncated"] is True


def test_dataframe_and_series_error_paths() -> None:
    class DataFrame:
        __module__ = "pandas.core.frame"

        @property
        def shape(self) -> tuple[int, int]:
            raise RuntimeError("shape fail")

        @property
        def columns(self) -> list[str]:
            raise RuntimeError("cols fail")

        def __len__(self) -> int:
            return 0

        def head(self, _rows: int) -> object:
            raise RuntimeError("head fail")

    class Series:
        __module__ = "pandas.core.series"

        @property
        def shape(self) -> tuple[int]:
            raise RuntimeError("shape fail")

        name = "bad"
        dtype = "int64"

        def head(self, _rows: int) -> object:
            raise RuntimeError("head fail")

    cfg = SnapshotConfig()
    df_summary = summarize_dataframe(DataFrame(), cfg)
    assert str(df_summary["shape"]).startswith("<error:")
    assert str(df_summary["columns"]).startswith("<error:")
    assert str(df_summary["head"]).startswith("<error:")

    series_summary = summarize_series(Series(), cfg)
    assert str(series_summary["shape"]).startswith("<error:")
    assert str(series_summary["head"]).startswith("<error:")


def test_jsonlike_dict_and_object_truncation_paths() -> None:
    cfg = SnapshotConfig(max_items=2)

    data = {1: "a", 2: "b", 3: "c"}
    out = to_jsonlike(data, cfg)
    assert isinstance(out, dict)
    assert "...[TRUNC]" in out

    class _NonDictDunder:
        @property
        def __dict__(self) -> int:  # type: ignore[override]
            return 123

        def __repr__(self) -> str:
            return "non-dict-repr"

    repr_out = to_jsonlike(_NonDictDunder(), cfg)
    assert repr_out == "non-dict-repr"

    class _BigObject:
        pass

    obj = _BigObject()
    for i in range(30):
        setattr(obj, f"k{i}", i)
    obj_out = to_jsonlike(obj, cfg)
    assert isinstance(obj_out, dict)
    assert "...[TRUNC]" in obj_out


def test_build_local_metadata_entry_handles_len_and_shape_errors() -> None:
    class _BadShape:
        def __iter__(self):
            raise RuntimeError("shape fail")

    class _ArrayLike:
        __module__ = "numpy"

        def __init__(self) -> None:
            self.shape = _BadShape()

        def __len__(self) -> int:
            raise RuntimeError("len fail")

    entry = _build_local_metadata_entry(_ArrayLike())
    assert "len" not in entry
    assert str(entry["shape"]).startswith("<error:")


def test_serialize_locals_bundle_handles_len_error_and_redaction_return_type(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    import llmdebug.serialize as serialize_module

    class _LenErrorMapping(dict[str, Any]):
        def __len__(self) -> int:
            raise RuntimeError("len fail")

    cfg = SnapshotConfig(max_items=2)
    bundle = serialize_locals_bundle(_LenErrorMapping({"a": 1, "b": 2}), cfg, [])
    assert bundle.truncated is False
    assert bundle.truncated_keys == 0

    monkeypatch.setattr(serialize_module, "redact_jsonlike", lambda *_a, **_k: ["not", "dict"])
    with pytest.raises(TypeError, match="expected dict"):
        serialize_locals_bundle({"x": 1}, cfg, [])
