# gRPC Server Implementation
import msgpack
import numpy as np
import grpc
from concurrent import futures
from neodynamics.interface import add_EnvironmentServiceServicer_to_server
from neodynamics.interface import EnvironmentType
from neodynamics.interface import EnvironmentService, SpacesResponse, ResetResponse, StepResponse, RenderResponse, Empty
from neodynamics.interface.utils import numpy_to_native, native_to_numpy, numpy_to_native_space, native_to_numpy_vec


class EnvironmentServicer(EnvironmentService):
    """gRPC servicer that wraps the GymEnvironment."""

    def __init__(self, environment_class):
        self.env = None
        self.environment_class = environment_class
        self.environment_type = None
        self.num_envs = None

    def Init(self, request, context):
        """Initialize the environment and return space information."""
        # Prepare initialization arguments
        init_args = {}
        if request.HasField('init_args'):
            init_args = msgpack.unpackb(request.init_args, raw=False)

        # Add render_mode to init_args if provided
        if request.HasField('render_mode'):
            init_args['render_mode'] = request.render_mode

        # Create the environment with all arguments
        self.env = self.environment_class(**init_args)

        # Create response with space information
        response = SpacesResponse()

        # Handle observation space (Dict space)
        for space_name, space in self.env.observation_space.spaces.items():
            space_proto = response.observation_space[space_name]
            numpy_to_native_space(space, space_proto)

        # Handle action space
        action_space = self.env.action_space
        numpy_to_native_space(action_space, response.action_space)

        if hasattr(self.env, 'num_envs'):
            self.num_envs = self.env.num_envs
            self.environment_type = EnvironmentType.VECTORIZED
        else:
            self.environment_type = EnvironmentType.STANDARD
            self.num_envs = 1

        response.num_envs = self.num_envs
        response.environment_type = self.environment_type

        return response

    def Reset(self, request, context):
        """Reset the environment and return the initial observation."""
        if self.env is None:
            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
            context.set_details("Environment not initialized. Call Init first.")
            return ResetResponse()

        # Extract seed and options if provided
        seed = None
        if request.HasField('seed'):
            seed = request.seed

        options = None
        if request.HasField('options'):
            options = msgpack.unpackb(request.options, raw=False)

        # Reset the environment
        obs, info = self.env.reset(seed=seed, options=options)

        # Convert numpy arrays to lists for serialization
        serializable_observation = self._get_serializable_observation(obs)

        # Serialize the observation and info
        response = ResetResponse(
            observation=msgpack.packb(serializable_observation, use_bin_type=True),
            info=msgpack.packb(info, use_bin_type=True)
        )

        return response

    def Step(self, request, context):
        """Take a step in the environment."""
        if self.env is None:
            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
            context.set_details("Environment not initialized. Call Init first.")
            return StepResponse()

        # Deserialize the action
        action = msgpack.unpackb(request.action, raw=False)

        # Convert lists back to numpy
        if self.environment_type == EnvironmentType.VECTORIZED:
            action = native_to_numpy_vec(action, self.env.action_space, self.num_envs)
        else:
            action = native_to_numpy(action, self.env.action_space)

        # Take a step in the environment
        obs, reward, terminated, truncated, info = self.env.step(action)

        # Convert numpy arrays to lists for serialization
        serializable_obs = self._get_serializable_observation(obs)

        # Create and return the response
        if self.environment_type == EnvironmentType.VECTORIZED:
            serializable_reward = reward.tolist()
            serializable_terminated = terminated.tolist()
            serializable_truncated = truncated.tolist()
        else:
            serializable_reward = float(reward)
            serializable_terminated = bool(terminated)
            serializable_truncated = bool(truncated)

        response = StepResponse(
            observation=msgpack.packb(serializable_obs, use_bin_type=True),
            reward=msgpack.packb(serializable_reward, use_bin_type=True),
            terminated=msgpack.packb(serializable_terminated, use_bin_type=True),
            truncated=msgpack.packb(serializable_truncated, use_bin_type=True),
            info=msgpack.packb(info, use_bin_type=True)
        )

        return response

    def Render(self, request, context):
        """Render the environment."""
        if self.env is None:
            context.set_code(grpc.StatusCode.FAILED_PRECONDITION)
            context.set_details("Environment not initialized. Call Init first.")
            return RenderResponse()

        # Get the render output directly
        render_output = self.env.render()

        # If it's a numpy array, directly serialize it
        if isinstance(render_output, np.ndarray):
            # Create a dict with array metadata and data for proper reconstruction
            array_data = {
                'shape': render_output.shape,
                'dtype': str(render_output.dtype),
                'data': render_output.tobytes(),
                'type': 'ndarray'
            }
            render_data = msgpack.packb(array_data, use_bin_type=True)
            return RenderResponse(render_data=render_data)
        else:
            # For non-array outputs, return empty data
            return RenderResponse(render_data=b'')

    def Close(self, request, context):
        """Close the environment."""
        if self.env is not None:
            self.env.close()
            self.env = None
            self.num_envs = None
            self.environment_type = None
        return Empty()

    def _get_serializable_observation(self, observation):
        if self.environment_type == EnvironmentType.VECTORIZED:
            return {key: value.tolist() for key, value in observation.items()}
        else:
            return {key: numpy_to_native(value, self.env.observation_space[key]) for key, value in observation.items()}

def create_environment_server(environment_class, port=50051):
    """Start the gRPC server."""
    environment_server = EnvironmentServicer(environment_class)
    server = grpc.server(futures.ThreadPoolExecutor(max_workers=10))
    add_EnvironmentServiceServicer_to_server(environment_server, server)
    server.add_insecure_port(f'[::]:{port}')
    server.start()
    print(f"Environment server started, listening on port {port}")
    server.wait_for_termination()