from typing import Any
import numpy as np

from google.protobuf import json_format, struct_pb2

from rlmesh.proto import gym_pb2 as pb

# ---------------- DType Maps ----------------
NP2DT = {
    np.uint8: pb.DType.UINT8,
    np.int32: pb.DType.INT32,
    np.int64: pb.DType.INT64,
    np.float32: pb.DType.FLOAT32,
    np.float64: pb.DType.FLOAT64,
}

DT2NP = {
    pb.DType.UINT8: np.uint8,
    pb.DType.INT32: np.int32,
    pb.DType.INT64: np.int64,
    pb.DType.FLOAT32: np.float32,
    pb.DType.FLOAT64: np.float64,
}


# ---------------- Helpers ----------------
def pack_mask(bits: np.ndarray) -> bytes:
    bits = np.asarray(bits, dtype=np.uint8).reshape(-1)
    n = int(bits.size)
    out = bytearray((n + 7) // 8)
    for i, b in enumerate(bits):
        if int(b):
            out[i >> 3] |= 1 << (i & 7)
    return bytes(out)


def unpack_mask(mask_bytes: bytes, n: int) -> np.ndarray:
    bits = np.frombuffer(mask_bytes, dtype=np.uint8)
    out = np.zeros(n, dtype=bool)
    for i in range(n):
        out[i] = bool(bits[i >> 3] & (1 << (i & 7)))
    return out


def struct_to_dict(s: struct_pb2.Struct | None) -> dict:
    if s is None:
        return {}
    return json_format.MessageToDict(s)


def dict_to_struct(d: dict[str, Any] | None) -> struct_pb2.Struct:
    return struct_pb2.Struct(
        fields={k: struct_pb2.Value(**_py_to_val(v)) for k, v in (d or {}).items()}  # pyright: ignore
    )


def _py_to_val(v):
    if v is None:
        return {"null_value": 0}

    # Python + NumPy scalars
    if isinstance(v, (bool, np.bool_)):
        return {"bool_value": bool(v)}
    if isinstance(v, (int, np.integer, float, np.floating)):
        return {"number_value": float(v)}

    # Strings
    if isinstance(v, str):
        return {"string_value": v}

    # NumPy arrays → ListValue (element-wise recursive conversion)
    if isinstance(v, np.ndarray):
        seq = v.tolist()
        return {
            "list_value": struct_pb2.ListValue(
                values=[struct_pb2.Value(**_py_to_val(x)) for x in seq]  # pyright: ignore
            )
        }

    # Python sequences → ListValue
    if isinstance(v, (list, tuple)):
        return {
            "list_value": struct_pb2.ListValue(
                values=[struct_pb2.Value(**_py_to_val(x)) for x in v]  # pyright: ignore
            )
        }

    # Dicts → Struct
    if isinstance(v, dict):
        return {"struct_value": dict_to_struct(v)}

    # Fallback: stringify
    return {"string_value": str(v)}


def unwrap_batched_info(batched_info: dict, index: int = 0) -> dict:
    """
    Extract a single environment's info dict (old-style) from a batched vector info.
    For each key in batched_info:
        - if it's a mask key (starts with "_"), skip it
        - if there's a mask "_key" and it's False for this index, skip the value
        - otherwise take batched_info[key][index]
    """
    if not isinstance(batched_info, dict):
        return {}

    out = {}
    for k, v in batched_info.items():
        if k.startswith("_"):
            continue
        mask_key = f"_{k}"
        if mask_key in batched_info:
            mask = batched_info[mask_key]
            if not mask[index]:
                continue
        if isinstance(v, np.ndarray) and v.ndim > 0:
            out[k] = v[index]
        else:
            out[k] = v
    return out
