# Copyright 2025 The RLinf Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Subprocess vectorized environment for Robocasa.

Based on metaworld/venv.py implementation, adapted for Robocasa/Robosuite environments.
"""

from multiprocessing import Pipe, connection
from multiprocessing.context import Process
from typing import Any, Callable, Optional, Union

import gymnasium as gym
import numpy as np

from rlinf.envs.venv import (
    BaseVectorEnv,
    CloudpickleWrapper,
    EnvWorker,
    ShArray,
    SubprocEnvWorker,
    SubprocVectorEnv,
    _setup_buf,
)


def _worker(
    parent: connection.Connection,
    p: connection.Connection,
    env_fn_wrapper: CloudpickleWrapper,
    obs_bufs: Optional[Union[dict, tuple, ShArray]] = None,
) -> None:
    """Worker function for robocasa subprocess environment.

    Based on metaworld's _worker function, adapted for robosuite environments.
    """

    def _encode_obs(
        obs: Union[dict, tuple, np.ndarray], buffer: Union[dict, tuple, ShArray]
    ) -> None:
        if isinstance(obs, np.ndarray) and isinstance(buffer, ShArray):
            buffer.save(obs)
        elif isinstance(obs, tuple) and isinstance(buffer, tuple):
            for o, b in zip(obs, buffer):
                _encode_obs(o, b)
        elif isinstance(obs, dict) and isinstance(buffer, dict):
            for k in obs.keys():
                _encode_obs(obs[k], buffer[k])
        return None

    def _check_success(env, env_return):
        success = env._check_success()
        env_return = list(env_return)
        info = env_return[-1]
        info["success"] = success
        env_return[-1] = info
        env_return = tuple(env_return)
        return env_return

    def get_ep_meta(env, env_return):
        ep_meta = env.get_ep_meta()
        env_return = list(env_return)
        info = env_return[-1]
        info["ep_meta"] = ep_meta
        env_return[-1] = info
        env_return = tuple(env_return)
        return env_return

    parent.close()
    env = env_fn_wrapper.data()
    try:
        while True:
            try:
                cmd, data = p.recv()
            except EOFError:  # the pipe has been closed
                p.close()
                break
            if cmd == "step":
                # Robosuite returns (obs, reward, done, info), not 5 values like gymnasium
                env_return = env.step(data)
                if obs_bufs is not None:
                    _encode_obs(env_return[0], obs_bufs)
                    env_return = (None, *env_return[1:])
                # RoboCasa step can't record success in info, _check_success() must be called
                if hasattr(env, "_check_success"):
                    env_return = _check_success(env, env_return)
                # call get_ep_meta() to get the RoboCasa env meta, includes prompt & layout_id, etcs
                if hasattr(env, "get_ep_meta"):
                    env_return = get_ep_meta(env, env_return)
                p.send(env_return)
            elif cmd == "reset":
                # Robosuite reset can return just obs or (obs, info)
                retval = env.reset(**data)
                reset_returns_info = (
                    isinstance(retval, (tuple, list))
                    and len(retval) == 2
                    and isinstance(retval[1], dict)
                )
                if reset_returns_info:
                    obs, info = retval
                else:
                    obs = retval
                    info = {}
                if obs_bufs is not None:
                    _encode_obs(obs, obs_bufs)
                    obs = None
                # call get_ep_meta() to get the RoboCasa env meta, includes prompt & layout_id, etcs
                if hasattr(env, "get_ep_meta"):
                    info = get_ep_meta(env, (info,))[-1]
                # return obs + info other than mere obs
                p.send((obs, info))
            elif cmd == "close":
                p.send(env.close())
                p.close()
                break
            elif cmd == "render":
                p.send(env.render(**data) if hasattr(env, "render") else None)
            elif cmd == "seed":
                if hasattr(env, "seed"):
                    p.send(env.seed(data))
                else:
                    env.reset(seed=data)
                    p.send(None)
            elif cmd == "getattr":
                p.send(getattr(env, data) if hasattr(env, data) else None)
            elif cmd == "setattr":
                setattr(env.unwrapped, data["key"], data["value"])
            else:
                p.close()
                raise NotImplementedError(f"Unknown command: {cmd}")
    except KeyboardInterrupt:
        p.close()


class RobocasaSubprocEnvWorker(SubprocEnvWorker):
    """Subprocess environment worker for Robocasa.

    Based on metaworld's ReconfigureSubprocEnvWorker, but without the reconfigure
    functionality since robocasa doesn't need it.
    """

    def __init__(self, env_fn: Callable[[], gym.Env], share_memory: bool = False):
        self.parent_remote, self.child_remote = Pipe()
        self.share_memory = share_memory
        self.buffer: Optional[Union[dict, tuple, ShArray]] = None
        if self.share_memory:
            dummy = env_fn()
            obs_space = dummy.observation_space
            dummy.close()
            del dummy
            self.buffer = _setup_buf(obs_space)
        args = (
            self.parent_remote,
            self.child_remote,
            CloudpickleWrapper(env_fn),
            self.buffer,
        )
        # Use our custom _worker function
        self.process = Process(target=_worker, args=args, daemon=True)
        self.process.start()
        self.child_remote.close()
        EnvWorker.__init__(self, env_fn)


class RobocasaSubprocEnv(SubprocVectorEnv):
    """Subprocess vectorized environment for Robocasa/Robosuite.

    Based on metaworld's ReconfigureSubprocEnv, adapted for robocasa environments.
    Uses subprocess isolation to avoid OpenGL context sharing issues in MuJoCo.
    """

    def __init__(self, env_fns: list[Callable[[], gym.Env]], **kwargs: Any) -> None:
        def worker_fn(fn: Callable[[], gym.Env]) -> RobocasaSubprocEnvWorker:
            # Use our custom worker with shared memory disabled
            # Robosuite dict observations work better without shared memory
            return RobocasaSubprocEnvWorker(fn, share_memory=False)

        BaseVectorEnv.__init__(self, env_fns, worker_fn, **kwargs)
