"""JSON-decoded tagged dicts → Python / NumPy / Pandas objects."""

from __future__ import annotations

import datetime
import math
from typing import Any

import numpy as np
import pandas as pd

from r_bridge.type_hints import TypeHint

_CUSTOM: dict[str, Any] = {}


def register_deserializer(r_tag: str, fn) -> None:
    _CUSTOM[r_tag] = fn


def deserialize(obj: Any, hint: TypeHint | None = None) -> Any:
    """Convert a JSON-decoded value from the R bridge into a Python object."""
    if hint == TypeHint.RAW:
        return obj

    if not isinstance(obj, dict) or "__type__" not in obj:
        # Plain JSON scalar or already-decoded value
        return obj

    tag = obj["__type__"]

    if tag in _CUSTOM:
        return _CUSTOM[tag](obj)

    dispatch = {
        "integer": _decode_integer,
        "numeric_scalar": lambda obj, hint: float(obj["value"]),
        "nan": _decode_nan,
        "inf": _decode_inf,
        "complex": _decode_complex,
        "posixct": _decode_posixct,
        "date": _decode_date,
        "na": _decode_na,
        "numeric_vector": _decode_numeric_vector,
        "integer_vector": _decode_integer_vector,
        "logical_vector": _decode_logical_vector,
        "character_vector": _decode_character_vector,
        "complex_vector": _decode_complex_vector,
        "matrix": _decode_matrix,
        "array": _decode_array,
        "dataframe": _decode_dataframe,
        "factor": _decode_factor,
        "named_list": _decode_named_list,
        "list": _decode_list,
    }

    fn = dispatch.get(tag)
    if fn is None:
        raise ValueError(f"r_bridge: unknown R type tag {tag!r}")
    return fn(obj, hint)


# ---------------------------------------------------------------------------
# Scalar decoders
# ---------------------------------------------------------------------------

def _decode_integer(obj, hint):
    return int(obj["value"])

def _decode_nan(obj, hint):
    return float("nan")

def _decode_inf(obj, hint):
    return math.copysign(math.inf, obj["sign"])

def _decode_complex(obj, hint):
    return complex(obj["r"], obj["i"])

def _decode_posixct(obj, hint):
    return datetime.datetime.fromisoformat(obj["iso8601"])

def _decode_date(obj, hint):
    return datetime.date.fromisoformat(obj["iso8601"])

def _decode_na(obj, hint):
    na_type = obj.get("na_type", "real")
    if na_type in ("real", "complex"):
        return float("nan")
    if na_type == "integer":
        return pd.NA
    # logical / character
    return None


# ---------------------------------------------------------------------------
# Vector decoders
# ---------------------------------------------------------------------------

def _decode_element(val: Any) -> Any:
    """Recursively decode a single element inside a vector."""
    if isinstance(val, dict) and "__type__" in val:
        return deserialize(val)
    return val


def _decode_numeric_vector(obj, hint):
    raw = [_decode_element(v) for v in obj["value"]]
    if hint == TypeHint.SCALAR and len(raw) == 1:
        return float(raw[0]) if not isinstance(raw[0], float) else raw[0]
    if hint == TypeHint.LIST:
        return raw
    arr = np.empty(len(raw), dtype=np.float64)
    for i, v in enumerate(raw):
        arr[i] = v if not (isinstance(v, float) and math.isnan(v)) else float("nan")
    return arr


def _decode_integer_vector(obj, hint):
    raw = [_decode_element(v) for v in obj["value"]]
    if hint == TypeHint.SCALAR and len(raw) == 1:
        return int(raw[0])
    if hint == TypeHint.LIST:
        return [int(v) for v in raw]
    return np.array(raw, dtype=np.int32)


def _decode_logical_vector(obj, hint):
    raw = [bool(v) for v in obj["value"]]
    if hint == TypeHint.SCALAR and len(raw) == 1:
        return raw[0]
    if hint == TypeHint.LIST:
        return raw
    return np.array(raw, dtype=bool)


def _decode_character_vector(obj, hint):
    raw = obj["value"]
    if hint == TypeHint.SCALAR and len(raw) == 1:
        return raw[0]
    return list(raw)


def _decode_complex_vector(obj, hint):
    raw = [_decode_element(v) for v in obj["value"]]
    if hint == TypeHint.SCALAR and len(raw) == 1:
        return complex(raw[0])
    return np.array(raw, dtype=np.complex128)


# ---------------------------------------------------------------------------
# Structured types
# ---------------------------------------------------------------------------

def _decode_matrix(obj, hint):
    data = [_decode_element(v) for v in obj["data"]]
    nrow, ncol = obj["nrow"], obj["ncol"]
    byrow = obj.get("byrow", True)
    arr = np.array(data, dtype=np.float64)
    return arr.reshape((nrow, ncol)) if byrow else arr.reshape((nrow, ncol), order="F")


def _decode_array(obj, hint):
    data = [_decode_element(v) for v in obj["data"]]
    shape = obj["shape"]
    return np.array(data).reshape(shape)


def _decode_dataframe(obj, hint):
    columns = {}
    for col_name, col_payload in obj["columns"].items():
        decoded = deserialize(col_payload, hint=TypeHint.LIST)
        columns[col_name] = decoded
    return pd.DataFrame(columns)


def _decode_factor(obj, hint):
    levels = obj["levels"]
    codes = obj["codes"]  # 1-based, -1 = NA
    values = [levels[c - 1] if c >= 1 else None for c in codes]
    return pd.Categorical(values, categories=levels, ordered=bool(obj.get("ordered", False)))


def _decode_named_list(obj, hint):
    return {k: deserialize(v) for k, v in obj["value"].items()}


def _decode_list(obj, hint):
    return [deserialize(v) for v in obj["value"]]
