"""Data storage operations."""
from __future__ import annotations

from pathlib import Path
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from .client import Client

_TRANSFER_TIMEOUT = 300.0


class Storage:
    def __init__(self, client: Client):
        self._client = client

    def list_buckets(self) -> list[dict]:
        """List available storage buckets."""
        return self._client.get("/v1/data/buckets")

    def list(self, prefix: str = "", bucket: str = "") -> list[dict]:
        """List objects, optionally filtered by prefix."""
        params: dict[str, str] = {}
        if prefix:
            params["prefix"] = prefix
        if bucket:
            return self._client.get(f"/v1/data/objects/{bucket}", params=params)
        return self._client.get("/v1/data/objects", params=params)

    def upload(self, key: str, local_path: str) -> dict:
        """Upload a local file.

        Uses ``multipart/form-data`` -- temporarily removes the JSON
        content-type header so ``requests`` sets the boundary correctly.
        """
        path = Path(local_path)
        if not path.is_file():
            raise FileNotFoundError(f"Local file not found: {local_path}")

        # Temporarily drop Content-Type so requests sets multipart boundary
        ct = self._client._session.headers.pop("Content-Type", None)
        try:
            with open(path, "rb") as fh:
                return self._client.post(
                    "/v1/data/objects",
                    files={"file": (path.name, fh)},
                    params={"key": key},
                    timeout=_TRANSFER_TIMEOUT,
                )
        finally:
            if ct:
                self._client._session.headers["Content-Type"] = ct

    def download(self, key: str, local_path: str) -> None:
        """Download a remote object to a local file."""
        url = f"{self._client.base_url}/v1/data/objects/{key}"
        resp = self._client._session.get(url, stream=True, timeout=_TRANSFER_TIMEOUT)
        if resp.status_code >= 400:
            # Let the standard error handler deal with it
            self._client._request("GET", f"/v1/data/objects/{key}")
        dest = Path(local_path)
        dest.parent.mkdir(parents=True, exist_ok=True)
        with open(dest, "wb") as fh:
            for chunk in resp.iter_content(chunk_size=8192):
                fh.write(chunk)

    def delete(self, key: str) -> None:
        """Delete a remote object."""
        self._client.delete(f"/v1/data/objects/{key}")

    def usage(self) -> dict:
        """Get storage usage statistics."""
        return self._client.get("/v1/data/usage")

    def quota(self) -> dict:
        """Get storage quota."""
        return self._client.get("/v1/data/quota")

    def datasets(self) -> list[dict]:
        """List registered datasets."""
        return self._client.get("/v1/data/datasets")
