from typing import Any

from yera.models.llm_context import llm_context


class Workspace:
    """Workspace for sharing data between agents within the same LLM context.

    Provides dict-like API for accessing workspace variables:
    - Use `workspace[key]` to access required keys (raises KeyError if missing)
    - Use `workspace.get(key, default)` for optional keys (returns default if missing)
    - Use `key in workspace` to check if a key exists
    - Use `workspace[key] = value` or `workspace.set(key, value)` to set values
    """

    def get(self, key: str, default: Any = None) -> Any:
        """Get a value from workspace, returning default if key doesn't exist.

        Args:
            key: The workspace key to retrieve
            default: Value to return if key doesn't exist (default: None)

        Returns:
            The value associated with key, or default if key doesn't exist
        """
        llm = llm_context()
        return llm.workspace.variables.get(key, default)

    def set(self, key: str, value: Any) -> None:
        """Set a value in workspace.

        Args:
            key: The workspace key to set
            value: The value to store
        """
        llm = llm_context()
        llm.workspace.variables[key] = value

    def __getitem__(self, key: str) -> Any:
        """Get a value from workspace, raising KeyError if key doesn't exist.

        Args:
            key: The workspace key to retrieve

        Returns:
            The value associated with key

        Raises:
            KeyError: If the key doesn't exist in workspace
        """
        llm = llm_context()
        if key not in llm.workspace.variables:
            raise KeyError(f"Workspace key '{key}' not found")
        return llm.workspace.variables[key]

    def __setitem__(self, key: str, value: Any) -> None:
        """Set a value in workspace.

        Args:
            key: The workspace key to set
            value: The value to store
        """
        self.set(key, value)

    def __contains__(self, key: str) -> bool:
        """Check if a key exists in workspace.

        Args:
            key: The workspace key to check

        Returns:
            True if key exists, False otherwise
        """
        llm = llm_context()
        return key in llm.workspace.variables


workspace = Workspace()
__all__ = ["workspace"]
