"""YarnClient -- core HTTP client with auth and error handling."""
from __future__ import annotations

import os

import requests

from .exceptions import (
    AuthenticationError,
    InsufficientCredits,
    NotFoundError,
    PermissionError,
    QuotaExceeded,
    RateLimitError,
    ServiceUnavailable,
    YarnError,
)

DEFAULT_BASE_URL = "https://api.au.yarn.prosodylabs.com.au"

_STATUS_MAP: dict[int, type[YarnError]] = {
    401: AuthenticationError,
    402: InsufficientCredits,
    403: PermissionError,
    404: NotFoundError,
    413: QuotaExceeded,
    429: RateLimitError,
    502: ServiceUnavailable,
    503: ServiceUnavailable,
}

_STATUS_MESSAGES: dict[int, str] = {
    401: "Invalid API key",
    402: "Insufficient credit balance",
    403: "Permission denied",
    404: "Not found",
    413: "Quota exceeded",
    429: "Rate limited",
    502: "Service unavailable",
    503: "Service unavailable",
}


class Client:
    """Yarn API client.

    Args:
        api_key: API key (``yarn_...``). Falls back to ``YARN_API_KEY`` env var.
        base_url: Override the default API endpoint.
        timeout: Default request timeout in seconds.
    """

    def __init__(
        self,
        api_key: str | None = None,
        base_url: str | None = None,
        timeout: float = 30.0,
    ):
        self.api_key = api_key or os.environ.get("YARN_API_KEY", "")
        if not self.api_key:
            raise AuthenticationError(
                "API key required. Set YARN_API_KEY or pass api_key="
            )
        self.base_url = (
            base_url or os.environ.get("YARN_BASE_URL", DEFAULT_BASE_URL)
        ).rstrip("/")
        self.timeout = timeout
        self._session = requests.Session()
        self._session.headers.update(
            {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json",
                "User-Agent": "yarn-au/0.1.1",
            }
        )

        # Lazily-initialised sub-clients
        self._jobs: Jobs | None = None  # type: ignore[name-defined]  # noqa: F821
        self._sessions: Sessions | None = None  # type: ignore[name-defined]  # noqa: F821
        self._notebooks: Notebooks | None = None  # type: ignore[name-defined]  # noqa: F821
        self._storage: Storage | None = None  # type: ignore[name-defined]  # noqa: F821
        self._billing: Billing | None = None  # type: ignore[name-defined]  # noqa: F821
        self._secrets: Secrets | None = None  # type: ignore[name-defined]  # noqa: F821
        self._chat: Chat | None = None  # type: ignore[name-defined]  # noqa: F821
        self._embeddings: Embeddings | None = None  # type: ignore[name-defined]  # noqa: F821
        self._batch: Batch | None = None  # type: ignore[name-defined]  # noqa: F821

    # -- Sub-client accessors (lazy imports to avoid circular deps) ----------

    @property
    def jobs(self):
        """Job submission and management."""
        if self._jobs is None:
            from .jobs import Jobs

            self._jobs = Jobs(self)
        return self._jobs

    @property
    def sessions(self):
        """Interactive GPU sessions."""
        if self._sessions is None:
            from .sessions import Sessions

            self._sessions = Sessions(self)
        return self._sessions

    @property
    def notebooks(self):
        """JupyterLab lifecycle."""
        if self._notebooks is None:
            from .notebooks import Notebooks

            self._notebooks = Notebooks(self)
        return self._notebooks

    @property
    def storage(self):
        """Data storage operations."""
        if self._storage is None:
            from .storage import Storage

            self._storage = Storage(self)
        return self._storage

    @property
    def billing(self):
        """Billing and credit operations."""
        if self._billing is None:
            from .billing import Billing

            self._billing = Billing(self)
        return self._billing

    @property
    def secrets(self):
        """Research secret management."""
        if self._secrets is None:
            from .secrets import Secrets

            self._secrets = Secrets(self)
        return self._secrets

    @property
    def chat(self):
        """Chat completions (inference)."""
        if self._chat is None:
            from .chat import Chat

            self._chat = Chat(self)
        return self._chat

    @property
    def embeddings(self):
        """Vector embeddings."""
        if self._embeddings is None:
            from .embeddings import Embeddings

            self._embeddings = Embeddings(self)
        return self._embeddings

    @property
    def batch(self):
        """Batch inference — run a prompt across a dataset."""
        if self._batch is None:
            from .batch import Batch

            self._batch = Batch(self)
        return self._batch

    # -- Top-level queries ---------------------------------------------------

    def gpus(self) -> dict:
        """List available GPUs with status, pricing, and model compatibility.

        Returns a dict with keys ``gpus``, ``models``, and ``pricing``.
        """
        return self.get("/v1/gpus")

    # -- HTTP helpers --------------------------------------------------------

    def _request(self, method: str, path: str, **kwargs) -> dict:
        """Send a request and translate HTTP errors to typed exceptions."""
        url = f"{self.base_url}{path}"
        kwargs.setdefault("timeout", self.timeout)
        resp = self._session.request(method, url, **kwargs)

        if resp.status_code >= 400:
            exc_cls = _STATUS_MAP.get(resp.status_code, YarnError)
            default_msg = _STATUS_MESSAGES.get(
                resp.status_code, f"API error {resp.status_code}"
            )
            try:
                body = resp.json()
                detail = body.get("detail", default_msg)
            except Exception:
                body = None
                detail = resp.text or default_msg
            raise exc_cls(detail, status_code=resp.status_code, response=body)

        if resp.status_code == 204:
            return {}
        return resp.json()

    def get(self, path: str, **kwargs) -> dict:
        return self._request("GET", path, **kwargs)

    def post(self, path: str, **kwargs) -> dict:
        return self._request("POST", path, **kwargs)

    def put(self, path: str, **kwargs) -> dict:
        return self._request("PUT", path, **kwargs)

    def delete(self, path: str, **kwargs) -> dict:
        return self._request("DELETE", path, **kwargs)
