"""Yarn Python SDK."""
from .client import Client
from .exceptions import (
    AuthenticationError,
    InsufficientCredits,
    NotFoundError,
    PermissionError,
    QuotaExceeded,
    RateLimitError,
    ServiceUnavailable,
    YarnError,
)
from .sessions import Session

__version__ = "0.1.0"


def session(
    name: str = "default",
    *,
    gpu: str = "rtx-4090",
    gpu_count: int = 1,
    api_key: str | None = None,
    **kwargs,
) -> Session:
    """Get a GPU session as a device-like context manager.

    Convenience wrapper — creates a Client and session in one call::

        import yarn

        with yarn.session(gpu="rtx4090") as device:
            model = MyModel().to(device)
            device.log({"epoch": 1, "loss": 0.42})
            device.save("model.pt", model.state_dict())

    Args:
        name: Session name (default: "default").
        gpu: GPU type (default: "rtx-4090").
        gpu_count: Number of GPUs (default: 1).
        api_key: API key. Falls back to YARN_API_KEY env var.
        **kwargs: Additional args passed to Sessions.session().
    """
    client = Client(api_key=api_key)
    return client.sessions.session(name, gpu=gpu, gpu_count=gpu_count, **kwargs)


__all__ = [
    "Client",
    "Session",
    "session",
    "YarnError",
    "AuthenticationError",
    "PermissionError",
    "InsufficientCredits",
    "NotFoundError",
    "QuotaExceeded",
    "RateLimitError",
    "ServiceUnavailable",
]
