from typing import Optional

import asyncio
import grpc
import numpy as np

from rlmesh.gym import (
    tensor_from_ndarray,
    ndarray_from_tensor,
)
from rlmesh.gym.utils import (
    struct_to_dict,
    unpack_mask,
)
from rlmesh.proto import env_pb2_grpc, env_pb2


class AsyncEnvTransport:
    def __init__(self, addr: str, metadata: list):
        self._addr = addr
        self._metadata = metadata

        self._chan: grpc.aio.Channel | None = None
        self._stub: env_pb2_grpc.EnvServiceStub | None = None
        self._stream: grpc.aio.StreamStreamCall | None = None
        self._ready = asyncio.Event()
        self._lock = asyncio.Lock()

    async def start(self):
        self._chan = grpc.aio.insecure_channel(self._addr)
        self._stub = env_pb2_grpc.EnvServiceStub(self._chan)
        self._stream = self._stub.Connect(metadata=self._metadata)
        self._ready.set()

    async def _read_and_check_msg(
        self, expected: Optional[str] = None
    ) -> env_pb2.ServerMsg:
        assert self._stream is not None
        msg = await self._stream.read()

        if msg is None:
            raise RuntimeError("stream closed on reset")

        assert isinstance(msg, env_pb2.ServerMsg)
        if msg.WhichOneof("kind") == "error":
            raise RuntimeError(
                f"recieved an error from the server: {msg.error.message}"
            )

        if expected is not None and msg.WhichOneof("kind") != expected:
            raise RuntimeError(f"expected {expected}, got {msg.WhichOneof('kind')}")

        return msg

    async def close(self):
        try:
            if self._stream is not None:
                await self._stream.write(env_pb2.ClientMsg(close=env_pb2.CloseReq()))
                await self._stream.done_writing()
        finally:
            if self._chan is not None:
                await self._chan.close()

    async def reset(
        self, *, seed: int | None, options: dict | None = None
    ) -> tuple[np.ndarray, dict]:
        await self._ready.wait()
        async with self._lock:
            assert self._stream is not None
            seeds = [] if seed is None else [int(seed)]
            reset = env_pb2.ResetReq(seeds=seeds)

            await self._stream.write(env_pb2.ClientMsg(reset=reset))
            msg = await self._read_and_check_msg("reset_ok")

            r = msg.reset_ok
            obs = ndarray_from_tensor(r.obs)
            infos = struct_to_dict(r.infos)

            return obs, infos

    async def step(
        self, actions: np.ndarray
    ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, dict]:
        await self._ready.wait()
        async with self._lock:
            assert self._stream is not None
            actions_ten = tensor_from_ndarray(actions)
            step = env_pb2.StepReq(actions=actions_ten)

            await self._stream.write(env_pb2.ClientMsg(step=step))
            msg = await self._read_and_check_msg("step_ok")

            s = msg.step_ok
            obs = ndarray_from_tensor(s.obs)
            rewards = ndarray_from_tensor(s.rewards)
            num_envs = rewards.shape[0]

            terminated = unpack_mask(s.terminated_mask, num_envs)
            truncated = unpack_mask(s.truncated_mask, num_envs)
            infos = struct_to_dict(s.infos)

            return obs, rewards, terminated, truncated, infos

    async def render(self) -> np.ndarray | None:
        await self._ready.wait()
        async with self._lock:
            assert self._stream is not None
            await self._stream.write(env_pb2.ClientMsg(render=env_pb2.RenderReq()))
            msg = await self._read_and_check_msg("render_ok")
            frame = ndarray_from_tensor(msg.render_ok.frame)
            return frame
