from typing import Any
from google.protobuf import json_format

import gymnasium as gym

from rlmesh.gym.spaces_codec import serialize_space, deserialize_space
from rlmesh.proto import gym_pb2 as pb


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


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


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


# --------------- Serialize: gym.Env -> EnvSpec ----------------
def serialize_env(env: gym.Env) -> pb.EnvSpec:
    package: str

    entry_point = getattr(getattr(env, "spec", None), "entry_point", None)
    if isinstance(entry_point, str):
        package = entry_point.split(".")[0]
    else:
        mod = getattr(env.__class__, "__module__", "") or ""
        package = mod.split(".")[0] if mod else "unknown"

    gym_id = (
        getattr(getattr(env, "spec", None), "id", None)
        or getattr(env, "spec_id", None)
        or ""
    )

    md = getattr(env, "metadata", {}) or {}
    render_modes = list(md.get("render_modes", []))
    render_fps = int(md.get("render_fps", 30))

    try:
        return pb.EnvSpec(
            package=package,
            gym_id=str(gym_id),
            spaces=pb.GymSpaces(
                obs=serialize_space(env.observation_space),
                action=serialize_space(env.action_space),
            ),
            metadata=pb.GymMetadata(
                render_modes=render_modes,
                render_fps=render_fps,
            ),
        )
    except Exception as e:
        raise SerializeError(f"Failed to serialize {gym_id} from {package}: {e}") from e


# --------------- Deserialize: EnvSpec -> gym.Env ---------------
def deserialize_env(env: pb.EnvSpec, *, validate: bool = True) -> gym.Env:
    observation_space = deserialize_space(env.spaces.obs, validate=validate)
    action_space = deserialize_space(env.spaces.action, validate=validate)

    class EnvShell(gym.Env):
        metadata: dict[str, Any] = {
            "render_modes": list(env.metadata.render_modes),
            "render_fps": int(env.metadata.render_fps or 30),
        }

        def __init__(self):
            super().__init__()
            self.observation_space = observation_space
            self.action_space = action_space

        def reset(
            self,
            *,
            seed: int | None = None,
            options: dict[str, Any] | None = None,
        ):
            super().reset(seed=seed)
            obs = self.observation_space.sample()
            return obs, {}

        def step(self, action):
            assert self.action_space.contains(action), "Action outside action_space"
            obs = self.observation_space.sample()
            return obs, 0, False, False, {}

        def render(self):
            return None

        def close(self):
            pass

    return EnvShell()


# ---------------- Helpers ----------------
def env_spec_to_json(e: pb.EnvSpec) -> str:
    return json_format.MessageToJson(e)


def env_spec_from_json(e: str) -> pb.EnvSpec:
    spec = pb.EnvSpec()
    json_format.Parse(e, spec)
    return spec
