import math
import datetime
import numpy as np
import pandas as pd
import pytest
from r_bridge.serializer import serialize, register_serializer, _CUSTOM


# --- Basic scalars ---

def test_none():
    assert serialize(None) is None

def test_bool():
    assert serialize(True) is True
    assert serialize(False) is False

def test_int():
    result = serialize(42)
    assert result == {"__type__": "integer", "value": 42}

def test_float():
    assert serialize(3.14) == pytest.approx(3.14)

def test_float_nan():
    assert serialize(float("nan")) == {"__type__": "nan"}

def test_float_inf():
    assert serialize(float("inf")) == {"__type__": "inf", "sign": 1}
    assert serialize(float("-inf")) == {"__type__": "inf", "sign": -1}

def test_str():
    assert serialize("hello") == "hello"

def test_complex():
    assert serialize(1 + 2j) == {"__type__": "complex", "r": 1.0, "i": 2.0}


# --- Collections ---

def test_homogeneous_int_list():
    result = serialize([1, 2, 3])
    assert result == {"__type__": "integer_vector", "value": [1, 2, 3]}

def test_homogeneous_float_list():
    result = serialize([1.0, 2.0])
    assert result == {"__type__": "numeric_vector", "value": [1.0, 2.0]}

def test_homogeneous_str_list():
    result = serialize(["a", "b"])
    assert result == {"__type__": "character_vector", "value": ["a", "b"]}

def test_homogeneous_bool_list():
    result = serialize([True, False])
    assert result == {"__type__": "logical_vector", "value": [True, False]}

def test_heterogeneous_list():
    result = serialize([1, "a"])
    assert result == {"__type__": "list", "value": [{"__type__": "integer", "value": 1}, "a"]}

def test_empty_list():
    result = serialize([])
    assert result == {"__type__": "list", "value": []}

def test_dict():
    result = serialize({"a": 1, "b": "x"})
    assert result == {"__type__": "named_list", "value": {"a": {"__type__": "integer", "value": 1}, "b": "x"}}


# --- NumPy ---

def test_numpy_1d_float():
    arr = np.array([1.0, 2.0, 3.0])
    result = serialize(arr)
    assert result == {"__type__": "numeric_vector", "value": [1.0, 2.0, 3.0]}

def test_numpy_1d_int():
    arr = np.array([1, 2, 3], dtype=np.int32)
    result = serialize(arr)
    assert result == {"__type__": "integer_vector", "value": [1, 2, 3]}

def test_numpy_1d_bool():
    arr = np.array([True, False])
    result = serialize(arr)
    assert result == {"__type__": "logical_vector", "value": [True, False]}

def test_numpy_2d():
    arr = np.array([[1, 2], [3, 4]], dtype=np.float64)
    result = serialize(arr)
    assert result["__type__"] == "matrix"
    assert result["nrow"] == 2
    assert result["ncol"] == 2
    assert result["data"] == [1.0, 2.0, 3.0, 4.0]

def test_numpy_nan_in_array():
    arr = np.array([1.0, float("nan"), 3.0])
    result = serialize(arr)
    assert result["__type__"] == "numeric_vector"
    assert result["value"][1] == {"__type__": "nan"}


# --- Pandas ---

def test_dataframe():
    df = pd.DataFrame({"x": [1, 2], "y": [3.0, 4.0]})
    result = serialize(df)
    assert result["__type__"] == "dataframe"
    assert "x" in result["columns"]
    assert "y" in result["columns"]

def test_dataframe_with_na():
    df = pd.DataFrame({"x": pd.array([1, None], dtype=pd.Int64Dtype())})
    result = serialize(df)
    col = result["columns"]["x"]
    assert col["__type__"] == "integer_vector"
    assert col["value"][1] == {"__type__": "na", "na_type": "integer"}

def test_series_numeric():
    s = pd.Series([1.0, 2.0], name="v")
    result = serialize(s)
    assert result == {"__type__": "numeric_vector", "value": [1.0, 2.0]}

def test_categorical():
    s = pd.Categorical(["a", "b", "a"], categories=["a", "b"])
    result = serialize(s)
    assert result["__type__"] == "factor"
    assert result["levels"] == ["a", "b"]
    assert result["ordered"] is False


# --- Datetime ---

def test_datetime():
    dt = datetime.datetime(2026, 3, 21, 12, 0, 0)
    result = serialize(dt)
    assert result["__type__"] == "posixct"
    assert "2026-03-21" in result["iso8601"]

def test_date():
    d = datetime.date(2026, 3, 21)
    result = serialize(d)
    assert result == {"__type__": "date", "iso8601": "2026-03-21"}


# --- Unknown type ---

def test_unknown_type_raises():
    class Foo:
        pass
    with pytest.raises(TypeError):
        serialize(Foo())


# --- Custom serializer ---

def test_register_serializer_custom_dispatch():
    class MyObj:
        def __init__(self, v): self.v = v
    register_serializer(MyObj, lambda obj: {"__type__": "myobj", "value": obj.v})
    try:
        assert serialize(MyObj(99)) == {"__type__": "myobj", "value": 99}
    finally:
        del _CUSTOM[MyObj]


# --- numpy scalar types ---

def test_numpy_float64_nan():
    assert serialize(np.float64("nan")) == {"__type__": "nan"}

def test_numpy_float64_inf():
    assert serialize(np.float64("inf")) == {"__type__": "inf", "sign": 1}
    assert serialize(np.float64("-inf")) == {"__type__": "inf", "sign": -1}

def test_numpy_bool_scalar():
    assert serialize(np.bool_(True)) is True
    assert serialize(np.bool_(False)) is False

def test_numpy_complex_scalar():
    result = serialize(np.complex128(1 + 2j))
    assert result == {"__type__": "complex", "r": 1.0, "i": 2.0}


# --- N-D arrays ---

def test_numpy_3d_array():
    arr = np.zeros((2, 3, 4))
    result = serialize(arr)
    assert result["__type__"] == "array"
    assert result["shape"] == [2, 3, 4]
    assert len(result["data"]) == 24

def test_numpy_inf_in_array():
    arr = np.array([1.0, float("inf"), float("-inf")])
    result = serialize(arr)
    assert result["__type__"] == "numeric_vector"
    assert result["value"][1] == {"__type__": "inf", "sign": 1}
    assert result["value"][2] == {"__type__": "inf", "sign": -1}


# --- Series dtype coverage ---

def test_series_bool_values():
    s = pd.Series([True, False])
    result = serialize(s)
    assert result["__type__"] == "logical_vector"

def test_series_string():
    s = pd.Series(["a", "b", "c"])
    result = serialize(s)
    assert result["__type__"] == "character_vector"
    assert result["value"] == ["a", "b", "c"]

def test_series_categorical_column():
    s = pd.Series(pd.Categorical(["x", "y", "x"], categories=["x", "y"]))
    result = serialize(s)
    assert result["__type__"] == "factor"

def test_series_datetime_dtype_falls_back_to_list():
    s = pd.Series(pd.to_datetime(["2026-01-01", "2026-01-02"]))
    result = serialize(s)
    assert result["__type__"] == "list"


# --- Series NA element coverage ---

def test_series_float_na():
    s = pd.Series([1.0, float("nan"), 3.0])
    result = serialize(s)
    assert result["__type__"] == "numeric_vector"
    assert result["value"][1] == {"__type__": "nan"}

def test_series_bool_na():
    s = pd.Series([True, pd.NA], dtype=pd.BooleanDtype())
    result = serialize(s)
    assert result["__type__"] == "logical_vector"
    assert result["value"][1] == {"__type__": "na", "na_type": "logical"}

def test_series_string_with_none():
    s = pd.Series(["a", None, "c"])
    result = serialize(s)
    assert result["__type__"] == "character_vector"
    assert result["value"][1] == {"__type__": "na", "na_type": "character"}
