"""Python / NumPy / Pandas → JSON-serializable tagged dicts for the R bridge."""

from __future__ import annotations

import datetime
import math
from typing import Any

import numpy as np
import pandas as pd

_CUSTOM: dict[type, Any] = {}


def register_serializer(python_type: type, fn) -> None:
    _CUSTOM[python_type] = fn


def serialize(obj: Any) -> Any:
    """Convert a Python object to a JSON-serializable form understood by bridge_loop.R."""
    # Custom serializers first
    for t, fn in _CUSTOM.items():
        if isinstance(obj, t):
            return fn(obj)

    # None
    if obj is None:
        return None

    # bool must come before int (bool is subclass of int in Python)
    if isinstance(obj, bool):
        return obj

    if isinstance(obj, int):
        return {"__type__": "integer", "value": obj}

    if isinstance(obj, float):
        if math.isnan(obj):
            return {"__type__": "nan"}
        if math.isinf(obj):
            return {"__type__": "inf", "sign": 1 if obj > 0 else -1}
        return obj

    if isinstance(obj, complex):
        return {"__type__": "complex", "r": obj.real, "i": obj.imag}

    if isinstance(obj, str):
        return obj

    if isinstance(obj, datetime.datetime):
        return {"__type__": "posixct", "iso8601": obj.isoformat()}

    if isinstance(obj, datetime.date):
        return {"__type__": "date", "iso8601": obj.isoformat()}

    if isinstance(obj, pd.Categorical):
        return _serialize_categorical(obj)

    if isinstance(obj, pd.DataFrame):
        return _serialize_dataframe(obj)

    if isinstance(obj, pd.Series):
        return _serialize_series(obj)

    if isinstance(obj, np.ndarray):
        return _serialize_ndarray(obj)

    if isinstance(obj, dict):
        return {"__type__": "named_list", "value": {k: serialize(v) for k, v in obj.items()}}

    if isinstance(obj, (list, tuple)):
        return _serialize_list(list(obj))

    # numpy scalar types
    if isinstance(obj, np.integer):
        return {"__type__": "integer", "value": int(obj)}
    if isinstance(obj, np.floating):
        v = float(obj)
        if math.isnan(v):
            return {"__type__": "nan"}
        if math.isinf(v):
            return {"__type__": "inf", "sign": 1 if v > 0 else -1}
        return v
    if isinstance(obj, np.bool_):
        return bool(obj)
    if isinstance(obj, np.complexfloating):
        return {"__type__": "complex", "r": float(obj.real), "i": float(obj.imag)}

    raise TypeError(f"r_bridge: cannot serialize object of type {type(obj).__name__!r}")


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _serialize_scalar(val: Any) -> Any:
    """Serialize a single scalar value that may appear inside a vector."""
    if val is None or (isinstance(val, float) and math.isnan(val)):
        return {"__type__": "nan"}
    return serialize(val)


def _all_same_type(items: list, *types) -> bool:
    return all(isinstance(i, types) and not isinstance(i, bool) for i in items) if types != (bool,) \
        else all(isinstance(i, bool) for i in items)


def _serialize_list(items: list) -> dict:
    if not items:
        return {"__type__": "list", "value": []}
    if all(isinstance(i, bool) for i in items):
        return {"__type__": "logical_vector", "value": items}
    if all(isinstance(i, int) and not isinstance(i, bool) for i in items):
        return {"__type__": "integer_vector", "value": items}
    if all(isinstance(i, float) for i in items):
        return {"__type__": "numeric_vector", "value": [_serialize_scalar(i) for i in items]}
    if all(isinstance(i, str) for i in items):
        return {"__type__": "character_vector", "value": items}
    return {"__type__": "list", "value": [serialize(i) for i in items]}


_NUMPY_DTYPE_MAP = {
    "float": "numeric_vector",
    "int": "integer_vector",
    "uint": "integer_vector",
    "bool": "logical_vector",
    "complex": "complex_vector",
    "str": "character_vector",
    "object": "list",
}


def _numpy_r_type(dtype: np.dtype) -> str:
    kind = dtype.kind  # 'f', 'i', 'u', 'b', 'c', 'U', 'O'
    return {
        "f": "numeric_vector",
        "i": "integer_vector",
        "u": "integer_vector",
        "b": "logical_vector",
        "c": "complex_vector",
        "U": "character_vector",
        "O": "list",
    }.get(kind, "list")


def _serialize_ndarray(arr: np.ndarray) -> dict:
    if arr.ndim == 1:
        r_type = _numpy_r_type(arr.dtype)
        values = [_serialize_element(v, arr.dtype) for v in arr.tolist()]
        return {"__type__": r_type, "value": values}
    if arr.ndim == 2:
        nrow, ncol = arr.shape
        flat = arr.flatten(order="C").tolist()
        values = [_serialize_element(v, arr.dtype) for v in flat]
        return {"__type__": "matrix", "nrow": nrow, "ncol": ncol, "data": values, "byrow": True}
    # N-D: treat as flat list with shape metadata
    shape = list(arr.shape)
    flat = arr.flatten().tolist()
    values = [_serialize_element(v, arr.dtype) for v in flat]
    return {"__type__": "array", "shape": shape, "data": values}


def _serialize_element(val: Any, dtype: np.dtype) -> Any:
    """Serialize a single element from a numpy array, handling NaN/Inf."""
    if dtype.kind == "f" and math.isnan(val):
        return {"__type__": "nan"}
    if dtype.kind == "f" and math.isinf(val):
        return {"__type__": "inf", "sign": 1 if val > 0 else -1}
    if dtype.kind == "b":
        return bool(val)
    return val


def _serialize_dataframe(df: pd.DataFrame) -> dict:
    columns = {}
    for col in df.columns:
        columns[str(col)] = _serialize_series(df[col])
    return {"__type__": "dataframe", "columns": columns}


def _serialize_series(s: pd.Series) -> dict:
    dtype = s.dtype
    if isinstance(dtype, pd.CategoricalDtype):
        return _serialize_categorical(s.cat)

    r_type = _pandas_dtype_to_r_vector_type(dtype)
    values = []
    for val in s:
        values.append(_serialize_series_element(val, dtype))
    return {"__type__": r_type, "value": values}


def _pandas_dtype_to_r_vector_type(dtype) -> str:
    if pd.api.types.is_bool_dtype(dtype):
        return "logical_vector"
    if pd.api.types.is_integer_dtype(dtype):
        return "integer_vector"
    if pd.api.types.is_float_dtype(dtype):
        return "numeric_vector"
    if pd.api.types.is_string_dtype(dtype) or dtype == object:
        return "character_vector"
    return "list"


def _serialize_series_element(val: Any, dtype) -> Any:
    """Serialize a single element from a pandas Series, handling NA."""
    if pd.isna(val):
        if pd.api.types.is_integer_dtype(dtype):
            return {"__type__": "na", "na_type": "integer"}
        if pd.api.types.is_float_dtype(dtype):
            return {"__type__": "nan"}
        if pd.api.types.is_bool_dtype(dtype):
            return {"__type__": "na", "na_type": "logical"}
        return {"__type__": "na", "na_type": "character"}
    if isinstance(val, bool):
        return val
    if isinstance(val, float) and math.isnan(val):
        return {"__type__": "nan"}
    return serialize(val)


def _serialize_categorical(cat) -> dict:
    levels = list(cat.categories)
    ordered = bool(cat.ordered)
    # codes as plain ints (1-based to match R convention)
    codes = [int(c) + 1 if c >= 0 else -1 for c in cat.codes]
    return {"__type__": "factor", "levels": levels, "codes": codes, "ordered": ordered}
