from typing import Any

import gymnasium as gym
import gymnasium.envs.registration as reg
import numpy as np

from rlmesh.env.env_transport import AsyncEnvTransport
from rlmesh.gym import deserialize_space
from rlmesh.proto import gym_pb2
from rlmesh.utils import LoopRunner


class RemoteEnv(gym.Env):
    metadata: dict[str, Any]

    def __init__(
        self,
        env_spec: gym_pb2.EnvSpec,
        *,
        transport_addr: str,
        transport_metadata: list,
    ):
        super().__init__()

        self.observation_space = deserialize_space(env_spec.spaces.obs)
        self.action_space = deserialize_space(env_spec.spaces.action)
        self.metadata = {
            "render_modes": list(env_spec.metadata.render_modes),
            "render_fps": int(env_spec.metadata.render_fps or 30),
        }

        self.spec = reg.EnvSpec(id=env_spec.gym_id)

        self._runner = LoopRunner()
        self._tx = AsyncEnvTransport(transport_addr, transport_metadata)
        self._runner.run(self._tx.start())
        self._closed = False

        self._obs = None
        self._is_ale = env_spec.package == "ale_py"

    def reset(self, *, seed: int | None = None, options: dict | None = None):
        if self._closed:
            raise RuntimeError("env closed")

        obs_batch, batched_info = self._runner.run(
            self._tx.reset(seed=seed, options=options)
        )
        obs = obs_batch[0]
        info = {}

        self._obs = obs
        return obs, info

    def step(self, action):
        if self._closed:
            raise RuntimeError("env closed")
        if not self.action_space.contains(action):
            raise ValueError("action outside action_space")

        arr = np.expand_dims(np.asarray(action), axis=0)
        obs_batch, rewards, terms, truns, infos = self._runner.run(self._tx.step(arr))

        obs = obs_batch[0]
        r = float(rewards[0])
        term = bool(terms[0])
        trunc = bool(truns[0])
        info = {}

        self._obs = obs
        return obs, r, term, trunc, info

    def render(self):
        if self._closed:
            raise RuntimeError("env closed")

        # Special Case for ALE Environment
        # - since we know the obs is the rgb_array, we can skip sending it over twice
        if self._is_ale and self._obs is not None:
            return self._obs

        frame = self._runner.run(self._tx.render())
        return frame

    def close(self):
        if not self._closed:
            try:
                self._runner.run(self._tx.close())
            finally:
                self._runner.stop()
                self._closed = True
