# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

"""
Generic environment client that works with raw dictionaries.

This module provides a GenericEnvClient that doesn't require installing
environment-specific packages. It's useful for connecting to remote servers
without running any untrusted code locally.
"""

from typing import Any, Dict

from .client_types import StepResult
from .env_client import EnvClient


class GenericEnvClient(EnvClient[Dict[str, Any], Dict[str, Any], Dict[str, Any]]):
    """
    Environment client that works with raw dictionaries instead of typed classes.

    This client doesn't require installing environment-specific packages, making it
    ideal for:
    - Connecting to remote servers without installing their packages
    - Quick prototyping and testing
    - Environments where type safety isn't needed
    - Security-conscious scenarios where you don't want to run remote code

    The trade-off is that you lose type safety and IDE autocomplete for actions
    and observations. Instead of typed objects, you work with plain dictionaries.

    Example:
        >>> # Direct connection to a running server (no installation needed)
        >>> with GenericEnvClient(base_url="http://localhost:8000") as env:
        ...     result = env.reset()
        ...     result = env.step({"code": "print('hello')"})
        ...     print(result.observation)  # Dict[str, Any]
        ...     print(result.observation.get("output"))

        >>> # From local Docker image
        >>> env = GenericEnvClient.from_docker_image("coding-env:latest")
        >>> result = env.reset()
        >>> result = env.step({"code": "x = 1 + 2"})
        >>> env.close()

        >>> # From HuggingFace Hub (pulls Docker image, no pip install)
        >>> env = GenericEnvClient.from_env("user/my-env", use_docker=True)
        >>> result = env.reset()
        >>> env.close()

    Note:
        GenericEnvClient inherits `from_docker_image()` and `from_env()` from
        EnvClient, so you can use it with Docker containers and HuggingFace
        Spaces without any package installation.
    """

    def _step_payload(self, action: Dict[str, Any]) -> Dict[str, Any]:
        """
        Convert action to payload for the server.

        For GenericEnvClient, this handles both raw dictionaries and
        typed Action objects (Pydantic models). If a Pydantic model is
        passed, it will be converted to a dictionary using model_dump().

        Args:
            action: Action as a dictionary or Pydantic BaseModel

        Returns:
            The action as a dictionary for the server
        """
        # If it's already a dict, return as-is
        if isinstance(action, dict):
            return action

        # If it's a Pydantic model (Action subclass), convert to dict
        if hasattr(action, "model_dump"):
            return action.model_dump()

        # Fallback for other objects with __dict__
        if hasattr(action, "__dict__"):
            return vars(action)

        # Last resort: try to convert to dict
        return dict(action)

    def _parse_result(self, payload: Dict[str, Any]) -> StepResult[Dict[str, Any]]:
        """
        Parse server response into a StepResult.

        Extracts the observation, reward, and done fields from the
        server response.

        Args:
            payload: Response payload from the server

        Returns:
            StepResult with observation as a dictionary
        """
        return StepResult(
            observation=payload.get("observation", {}),
            reward=payload.get("reward"),
            done=payload.get("done", False),
        )

    def _parse_state(self, payload: Dict[str, Any]) -> Dict[str, Any]:
        """
        Parse state response from the server.

        For GenericEnvClient, this returns the payload as-is since
        we're working with dictionaries.

        Args:
            payload: State payload from the server

        Returns:
            The state as a dictionary
        """
        return payload


class GenericAction(Dict[str, Any]):
    """
    A dictionary subclass for creating actions when using GenericEnvClient.

    This provides a semantic wrapper around dictionaries to make code more
    readable when working with GenericEnvClient. It behaves exactly like a
    dict but signals intent that this is an action for an environment.

    Example:
        >>> # Without GenericAction (works fine)
        >>> env.step({"code": "print('hello')"})

        >>> # With GenericAction (more explicit)
        >>> action = GenericAction(code="print('hello')")
        >>> env.step(action)

        >>> # With multiple fields
        >>> action = GenericAction(code="x = 1", timeout=30, metadata={"tag": "test"})
        >>> env.step(action)

    Note:
        GenericAction is just a dict with a constructor that accepts keyword
        arguments. It's provided for symmetry with typed Action classes and
        to make code more readable.
    """

    def __init__(self, **kwargs: Any) -> None:
        """
        Create a GenericAction from keyword arguments.

        Args:
            **kwargs: Action fields as keyword arguments

        Example:
            >>> action = GenericAction(code="print(1)", timeout=30)
            >>> action["code"]
            'print(1)'
        """
        super().__init__(kwargs)

    def __repr__(self) -> str:
        """Return a readable representation."""
        items = ", ".join(f"{k}={v!r}" for k, v in self.items())
        return f"GenericAction({items})"