import numpy as np
import gymnasium as gym

from google.protobuf import json_format

from rlmesh.proto import gym_pb2 as pb
from rlmesh.gym.utils import NP2DT, DT2NP


# ---------------- Errors ----------------
class Error(Exception):
    """Top-level module error for space_codec."""


class SerializeError(Error):
    """Thrown if serialization to SpaceSpec fails."""


class DeserializeError(Error):
    """Thrown if unable to deserialize SpaceSpec into a gym.Space"""


# --------------- Serialize: gym.Space -> SpaceSpec ----------------
def serialize_space(space: gym.Space, *, _path: str = "$") -> pb.SpaceSpec:
    try:
        # Box
        if isinstance(space, gym.spaces.Box):
            dt = NP2DT.get(space.dtype.type, pb.DType.FLOAT32)
            shape = list(space.shape)

            lo = np.asarray(space.low, dtype=np.float64)
            hi = np.asarray(space.high, dtype=np.float64)

            # 1) unbounded
            if np.all(np.isneginf(lo)) and np.all(np.isposinf(hi)):
                box = pb.BoxSpec(unbounded=True)

            # 2) scalar: either provided as scalar OR arrays that are all-equal
            elif (lo.size == 1 and hi.size == 1) or (
                np.all(lo == lo.flat[0]) and np.all(hi == hi.flat[0])
            ):
                lo_s = float(lo.reshape(())) if lo.size == 1 else float(lo.flat[0])
                hi_s = float(hi.reshape(())) if hi.size == 1 else float(hi.flat[0])
                box = pb.BoxSpec(lo_scalar=lo_s, hi_scalar=hi_s)

            # 3) per-axis broadcast (rank-length vectors)
            elif lo.ndim == 1 and hi.ndim == 1 and lo.shape[0] == len(shape):
                # Provided as per-axis vectors already
                box = pb.BoxSpec(lo_broadcast=lo.tolist(), hi_broadcast=hi.tolist())

            else:
                # 4) elementwise or any other pattern we don't encode losslessly -> envelope
                lo_min = float(np.min(lo))
                hi_max = float(np.max(hi))
                box = pb.BoxSpec(lo_scalar=lo_min, hi_scalar=hi_max)

            return pb.SpaceSpec(kind=pb.SpaceKind.BOX, dtype=dt, shape=shape, box=box)

        # Discrete
        if isinstance(space, gym.spaces.Discrete):
            return pb.SpaceSpec(
                kind=pb.SpaceKind.DISCRETE,
                dtype=pb.DType.INT64,  # canonical
                discrete=pb.DiscreteSpec(
                    n=int(space.n), start=int(getattr(space, "start", 0))
                ),
            )

        # MultiBinary (int or sequence)
        if isinstance(space, gym.spaces.MultiBinary):
            shape = list(space.shape)
            n = int(np.prod(shape))
            return pb.SpaceSpec(
                kind=pb.SpaceKind.MULTIBINARY,
                dtype=pb.DType.UINT8,
                shape=shape,  # canonical: carry shape; n is convenience
                multi_binary=pb.MultiBinarySpec(n=n),
            )

        # MultiDiscrete (can be ndarray with axes) -> flatten + carry shape
        if isinstance(space, gym.spaces.MultiDiscrete):
            dt = NP2DT.get(space.dtype.type, pb.DType.INT64)
            nvec = np.asarray(space.nvec, dtype=np.int64)
            shp = list(nvec.shape)  # [] if 1D vector
            return pb.SpaceSpec(
                kind=pb.SpaceKind.MULTIDISCRETE,
                dtype=dt,
                shape=shp,  # shape of nvec
                multi_discrete=pb.MultiDiscreteSpec(nvec=nvec.reshape(-1).tolist()),
            )

        # Text
        if isinstance(space, gym.spaces.Text):
            return pb.SpaceSpec(
                kind=pb.SpaceKind.TEXT,
                text=pb.TextSpec(
                    min_length=int(space.min_length),
                    max_length=int(space.max_length),
                    charset=str(space.characters),
                ),
            )

        # Dict
        if isinstance(space, gym.spaces.Dict):
            keys = list(space.spaces.keys())
            subs = [
                serialize_space(space.spaces[k], _path=f"{_path}.{k}") for k in keys
            ]
            return pb.SpaceSpec(
                kind=pb.SpaceKind.DICT, dict=pb.DictSpec(keys=keys, spaces=subs)
            )

        # Tuple
        if isinstance(space, gym.spaces.Tuple):
            subs = [
                serialize_space(s, _path=f"{_path}[{i}]")
                for i, s in enumerate(space.spaces)
            ]
            return pb.SpaceSpec(
                kind=pb.SpaceKind.TUPLE, tuple=pb.TupleSpec(spaces=subs)
            )

        raise TypeError(f"Unsupported Gym Space type: {type(space)}")
    except Exception as e:
        raise SerializeError(f"Failed to serialize {type(space).__name__}: {e}") from e


# --------------- Deserialize: SpaceSpec -> gym.Space ---------------
def deserialize_space(
    spec: pb.SpaceSpec, *, _path: str = "$", validate: bool = True
) -> gym.Space:
    try:
        validate_space_spec(spec)

        k = spec.kind
        if k == pb.SpaceKind.BOX:
            dt = DT2NP.get(spec.dtype, np.float32)
            shape = tuple(int(x) for x in spec.shape)
            b = spec.box

            # 1) unbounded
            if b.unbounded:
                lo = np.full(shape, -np.inf, dtype=np.float64)
                hi = np.full(shape, np.inf, dtype=np.float64)

            # 2) scalar
            elif (b.lo_scalar or b.hi_scalar) or (
                b.lo_scalar == 0.0 and b.hi_scalar == 0.0
            ):
                lo = np.full(shape, float(b.lo_scalar), dtype=np.float64)
                hi = np.full(shape, float(b.hi_scalar), dtype=np.float64)

            # 3) per-axis broadcast (rank-length vectors)
            elif len(b.lo_broadcast) and len(b.hi_broadcast):
                lo_vec = np.asarray(b.lo_broadcast, dtype=np.float64)
                hi_vec = np.asarray(b.hi_broadcast, dtype=np.float64)
                if lo_vec.shape[0] != len(shape) or hi_vec.shape[0] != len(shape):
                    raise ValueError(
                        f"BoxSpec broadcast lengths must equal rank={len(shape)}"
                    )
                # Broadcast per axis
                lo = np.broadcast_to(
                    lo_vec.reshape([len(shape)] + [1] * (len(shape) - 1)), shape
                ).copy()
                hi = np.broadcast_to(
                    hi_vec.reshape([len(shape)] + [1] * (len(shape) - 1)), shape
                ).copy()

            else:
                # Fallback: treat as unbounded if no variant is set
                lo = np.full(shape, -np.inf, dtype=np.float64)
                hi = np.full(shape, np.inf, dtype=np.float64)

            return gym.spaces.Box(low=lo.astype(dt), high=hi.astype(dt), dtype=dt)

        if k == pb.SpaceKind.DISCRETE:
            n = int(spec.discrete.n)
            start = int(spec.discrete.start)
            if n <= 0:
                raise ValueError(f"Discrete.n must be > 0, got {n}")
            return gym.spaces.Discrete(n=n, start=start)

        if k == pb.SpaceKind.MULTIBINARY:
            shape = (
                tuple(int(x) for x in spec.shape)
                if len(spec.shape)
                else (int(spec.multi_binary.n),)
            )
            return gym.spaces.MultiBinary(shape)

        if k == pb.SpaceKind.MULTIDISCRETE:
            flat = np.array(spec.multi_discrete.nvec, dtype=np.int64)
            shp = tuple(int(x) for x in spec.shape)  # shape of nvec
            nvec = flat.reshape(shp) if len(shp) else flat
            if np.any(nvec <= 0):
                raise ValueError(f"MultiDiscrete.nvec must be > 0, got {nvec}")
            return gym.spaces.MultiDiscrete(nvec)

        if k == pb.SpaceKind.TEXT:
            ts = spec.text
            kwargs = {}
            if ts.min_length:
                kwargs["min_length"] = int(ts.min_length)
            if ts.max_length:
                kwargs["max_length"] = int(ts.max_length)
            if ts.charset:
                kwargs["charset"] = ts.charset
            return gym.spaces.Text(**kwargs)

        if k == pb.SpaceKind.DICT:
            keys = list(spec.dict.keys)
            subs = list(spec.dict.spaces)
            if len(keys) != len(subs):
                raise ValueError("DictSpec: keys and spaces length mismatch")
            return gym.spaces.Dict(
                {
                    k: deserialize_space(s, _path=f"{_path}.{k}")
                    for k, s in zip(keys, subs)
                }
            )

        if k == pb.SpaceKind.TUPLE:
            return gym.spaces.Tuple(
                tuple(
                    deserialize_space(s, _path=f"{_path}[{i}]")
                    for i, s in enumerate(spec.tuple.spaces)
                )
            )

        raise ValueError(f"Unsupported SpaceKind: {k}")
    except Exception as e:
        raise DeserializeError(f"Failed to deserialize: {e}") from e


# ---------------- Validators ----------------
def validate_space_spec(spec: pb.SpaceSpec) -> None:
    k = spec.kind

    if k == pb.SpaceKind.BOX:
        if not len(spec.shape):
            raise ValueError("Box.shape must be set")
        b = spec.box
        use_scalar = (b.lo_scalar or b.hi_scalar) and not (
            b.lo_scalar == 0.0 and b.hi_scalar == 0.0
        )
        active = (
            int(bool(b.unbounded))
            + int(bool(use_scalar))
            + int(bool(len(b.lo_broadcast) or len(b.hi_broadcast)))
        )
        if active != 1:
            raise ValueError(
                "BoxSpec must use exactly one of: unbounded | scalar | broadcast"
            )
        if use_scalar and not (
            (b.lo_scalar or b.lo_scalar == 0.0) and (b.hi_scalar or b.hi_scalar == 0.0)
        ):
            raise ValueError("BoxSpec scalar: both lo_scalar and hi_scalar must be set")
        if len(b.lo_broadcast) != len(b.hi_broadcast) and (
            len(b.lo_broadcast) or len(b.hi_broadcast)
        ):
            raise ValueError(
                "BoxSpec: lo_broadcast and hi_broadcast lengths must match"
            )
        if len(b.lo_broadcast) and len(b.lo_broadcast) != len(spec.shape):
            raise ValueError("BoxSpec: broadcast length must equal rank of shape")

    elif k == pb.SpaceKind.DISCRETE:
        if spec.discrete.n <= 0:
            raise ValueError("Discrete.n must be > 0")

    elif k == pb.SpaceKind.MULTIBINARY:
        if not (len(spec.shape) or spec.multi_binary.n):
            raise ValueError("MultiBinary must carry shape or n")
        if len(spec.shape) and np.prod(spec.shape) <= 0:
            raise ValueError("MultiBinary.shape must be positive")

    elif k == pb.SpaceKind.MULTIDISCRETE:
        if not len(spec.multi_discrete.nvec):
            raise ValueError("MultiDiscrete.nvec must not be empty")
        if any(v <= 0 for v in spec.multi_discrete.nvec):
            raise ValueError("MultiDiscrete.nvec values must be > 0")

    elif k == pb.SpaceKind.DICT:
        if len(spec.dict.keys) != len(spec.dict.spaces):
            raise ValueError("DictSpec keys/spaces length mismatch")
        seen = set()
        for k_ in spec.dict.keys:
            if k_ in seen:
                raise ValueError(f"DictSpec duplicate key: {k_}")
            seen.add(k_)
        for s in spec.dict.spaces:
            validate_space_spec(s)

    elif k == pb.SpaceKind.TUPLE:
        if not len(spec.tuple.spaces):
            raise ValueError("TupleSpec must contain at least one space")
        for s in spec.tuple.spaces:
            validate_space_spec(s)

    elif k == pb.SpaceKind.TEXT:
        ts = spec.text
        if ts.max_length and ts.min_length and ts.max_length < ts.min_length:
            raise ValueError("TextSpec: max_length must be >= min_length")


# ---------------- Helpers ----------------
def space_spec_to_json(spec: pb.SpaceSpec) -> str:
    return json_format.MessageToJson(spec)


def space_spec_from_json(s: str) -> pb.SpaceSpec:
    spec = pb.SpaceSpec()
    json_format.Parse(s, spec)
    return spec
