app/backend/fluksio/sdk/client.py

"""Talking to an engine: uploading declared flows, and running them.

Two halves of one job. :func:`sync` puts what the decorators declared into the
flow store; :class:`Client` submits runs and reads back what they produced, so
an experiment is started and inspected from the same script that defines it.
"""

from __future__ import annotations

import json
import os
import subprocess
import time
from collections.abc import Iterable
from pathlib import Path
from typing import Any

from fluksio.sdk import MARKER, Flow, SyncError

__all__ = [
    "Client",
    "RunHandle",
    "SyncReport",
    "config_path",
    "data_dir",
    "find_data_dir",
    "ignore_self",
    "login",
    "sync",
    "write_config",
]

API = "/api/v1"
#: A run is over when it reaches one of these.
DONE = frozenset({"ok", "error", "cancelled", "abandoned"})


#: What a project-local installation is called, beside `.venv` and `.git`.
DATA_DIR_NAME = ".fluksio"

#: The shared one, for a machine that wants a single engine rather than one
#: per repository. `fluksio serve --global` is how you ask for it.
GLOBAL_DATA_DIR = Path("~/.fluksio")


def find_data_dir(start: Path | None = None) -> Path | None:
    """The nearest project-local installation, walking up from ``start``.

    The same search `.git` and `.venv` get, and for the same reason: which
    installation you mean is a fact about where you are standing, not about
    which machine you are on. Several repositories on one device each keep
    their own flows, runs and token this way rather than sharing one.
    """
    directory = (start or Path.cwd()).resolve()
    for candidate in (directory, *directory.parents):
        local = candidate / DATA_DIR_NAME
        if local.is_dir():
            return local
    return None


def data_dir(start: Path | None = None) -> Path:
    """The installation this working directory belongs to."""
    found = find_data_dir(start)
    return found if found is not None else GLOBAL_DATA_DIR.expanduser()


def config_path(directory: Path | None = None) -> Path:
    """Where the token for an installation lives — beside the data it opens.

    Not a single file per machine: a token is for one engine, and with an
    installation per repository there is more than one. Keeping it inside the
    data directory means the client finds the credential for the engine whose
    directory it is standing in, without either having to be told.
    """
    return (directory or data_dir()) / "client.json"


def _legacy_config_path() -> Path:
    """Where `fluksio login` used to write, before installations were local."""
    base = os.environ.get("XDG_CONFIG_HOME") or str(Path.home() / ".config")
    return Path(base) / "fluksio" / "client.json"


def _read(path: Path) -> dict[str, str]:
    if not path.exists():
        return {}
    try:
        data = json.loads(path.read_text())
    except ValueError:
        return {}
    return data if isinstance(data, dict) else {}


def _stored() -> dict[str, str]:
    """The nearest credential: this project's, then the machine's."""
    for path in (config_path(), _legacy_config_path()):
        found = _read(path)
        if found:
            return found
    return {}


class ApiError(Exception):
    """The engine refused, with what it said."""

    def __init__(self, status: int, detail: Any) -> None:
        super().__init__(f"{status}: {detail}")
        self.status = status
        self.detail = detail


class Conflict(ApiError):
    """Someone else saved in between — the 409 the optimistic lock answers."""

    def __init__(self, detail: Any) -> None:
        super().__init__(409, detail)
        self.current_version = 0
        if isinstance(detail, dict):
            self.current_version = int(detail.get("current_version") or 0)


class Client:
    """An authenticated engine, addressed over its HTTP API."""

    def __init__(
        self, url: str = "", token: str = "", http: Any = None, timeout: float = 30.0
    ) -> None:
        stored = _stored()
        self.url = url or os.environ.get("FLUKSIO_URL") or stored.get("url") or ""
        self.token = (
            token or os.environ.get("FLUKSIO_TOKEN") or stored.get("token") or ""
        )
        if http is None:
            if not self.url:
                raise SyncError(
                    "No engine to talk to. Run `fluksio login --url http://…`, or "
                    "set FLUKSIO_URL and FLUKSIO_TOKEN."
                )
            import httpx

            http = httpx.Client(base_url=self.url, timeout=timeout)
        self.http = http
        if self.token:
            self.http.headers["Authorization"] = f"Bearer {self.token}"

    # -- plumbing ----------------------------------------------------------

    def _call(self, method: str, path: str, **kwargs: Any) -> Any:
        response = self.http.request(method, f"{API}{path}", **kwargs)
        if response.status_code == 409:
            raise Conflict(_detail(response))
        if response.status_code >= 400:
            raise ApiError(response.status_code, _detail(response))
        if response.status_code == 204 or not response.content:
            return None
        return response.json()

    # -- flows -------------------------------------------------------------

    def get_flow(self, name: str) -> dict[str, Any] | None:
        """The stored flow, draft included, or ``None`` if there is none."""
        try:
            result: dict[str, Any] = self._call("GET", f"/flows/{name}")
        except ApiError as exc:
            if exc.status == 404:
                return None
            raise
        return result

    def put_flow(self, document: dict[str, Any]) -> dict[str, Any]:
        result: dict[str, Any] = self._call(
            "PUT", f"/flows/{document['name']}", json=document
        )
        return result

    def get_source(self, flow: str, node: str) -> str:
        result = self._call("GET", f"/flows/{flow}/nodes/{node}/source")
        return str(result.get("code", ""))

    def put_source(self, flow: str, node: str, code: str) -> dict[str, Any]:
        result: dict[str, Any] = self._call(
            "PUT", f"/flows/{flow}/nodes/{node}/source", json={"code": code}
        )
        return result

    def publish(self, name: str, version: int) -> dict[str, Any]:
        result: dict[str, Any] = self._call(
            "POST", f"/flows/{name}/publish", json={"version": version}
        )
        return result

    def refresh_modules(self) -> None:
        """Retire the engine's workers, so the next run imports the code as it is."""
        self._call("POST", "/modules/refresh")

    # -- runs --------------------------------------------------------------

    def submit(
        self,
        flow: str,
        params: dict[str, Any] | None = None,
        seed: int | None = None,
        no_cache: bool = False,
    ) -> RunHandle:
        row = self._call(
            "POST",
            f"/runs/flows/{flow}",
            json={"params": params or {}, "seed": seed, "no_cache": no_cache},
        )
        return RunHandle(self, row["id"], row)

    def sweep(
        self,
        flow: str,
        entries: list[dict[str, Any]],
        no_cache: bool = False,
    ) -> list[RunHandle]:
        """Many runs of one flow at once. The caller decides what varies."""
        rows = self._call(
            "POST",
            f"/runs/flows/{flow}/sweep",
            json={"runs": entries, "no_cache": no_cache},
        )
        return [RunHandle(self, row["id"], row) for row in rows]

    def run(self, run_id: str) -> dict[str, Any]:
        result: dict[str, Any] = self._call("GET", f"/runs/{run_id}")
        return result

    def runs(self, flow: str = "", limit: int = 20, **filters: Any) -> Any:
        query = {"limit": limit, **filters}
        if flow:
            query["flow"] = flow
        return self._call("GET", "/runs", params=query)

    def metrics(self, run_id: str, name: str = "", stride: int = 1) -> Any:
        query: dict[str, Any] = {"stride": stride}
        if name:
            query["name"] = name
        return self._call("GET", f"/runs/{run_id}/metrics", params=query)

    def compare(self, ids: Iterable[str], metric: str) -> Any:
        return self._call(
            "GET",
            "/runs/series/compare",
            params={"ids": ",".join(ids), "metric": metric},
        )

    def cancel(self, run_id: str) -> Any:
        return self._call("POST", f"/runs/{run_id}/cancel")

    def download(self, digest: str) -> bytes:
        response = self.http.request("GET", f"{API}/artifacts/{digest}")
        if response.status_code >= 400:
            raise ApiError(response.status_code, _detail(response))
        return bytes(response.content)


def _detail(response: Any) -> Any:
    try:
        body = response.json()
    except ValueError:
        return response.text
    return body.get("detail", body) if isinstance(body, dict) else body


class RunHandle:
    """One run, and the answers it accumulates."""

    def __init__(self, client: Client, run_id: str, row: dict[str, Any]) -> None:
        self.client = client
        self.id = run_id
        self._row = row

    def refresh(self) -> RunHandle:
        self._row = self.client.run(self.id)
        return self

    @property
    def status(self) -> str:
        return str(self._row.get("status", ""))

    @property
    def done(self) -> bool:
        return self.status in DONE

    @property
    def result(self) -> dict[str, Any]:
        """What the flow's outputs held when it finished."""
        result = self._row.get("result")
        return result if isinstance(result, dict) else {}

    @property
    def artifacts(self) -> list[dict[str, Any]]:
        rows = self._row.get("artifacts")
        return rows if isinstance(rows, list) else []

    def wait(self, timeout: float = 0.0, poll: float = 1.0) -> RunHandle:
        """Block until the run is over, or ``timeout`` seconds have passed."""
        deadline = time.monotonic() + timeout if timeout else 0.0
        while True:
            self.refresh()
            if self.done:
                return self
            if deadline and time.monotonic() > deadline:
                raise TimeoutError(f"run {self.id} is still {self.status}")
            time.sleep(poll)

    def metrics(self, name: str = "", stride: int = 1) -> list[dict[str, Any]]:
        """A streamed port's whole series — a run's metrics are its outputs."""
        points: list[dict[str, Any]] = self.client.metrics(self.id, name, stride)
        return points

    def download(self, name: str) -> bytes:
        """The bytes of an artifact this run produced."""
        for row in self.artifacts:
            if row.get("name") == name:
                return self.client.download(str(row["digest"]))
        have = ", ".join(str(row.get("name")) for row in self.artifacts) or "none"
        raise KeyError(f"run {self.id} has no artifact '{name}' (it has {have})")

    def __getitem__(self, key: str) -> Any:
        return self._row[key]

    def __repr__(self) -> str:
        return f"RunHandle({self.id!r}, status={self.status!r})"


def ignore_self(directory: Path) -> None:
    """Keep an installation out of the repository it sits in.

    It holds a token and a database, neither of which belongs in anybody's
    history. A `.gitignore` of `*` inside the directory ignores it from within,
    so nothing has to be added to the project's own — the same thing `uv` does
    for the venv it builds.
    """
    marker = directory / ".gitignore"
    if not marker.exists():
        directory.mkdir(parents=True, exist_ok=True)
        marker.write_text("# A Fluksio installation: a database, and a token.\n*\n")


def write_config(url: str, token: str, directory: Path | None = None) -> Path:
    """Store the credential for an engine, readable only by its owner."""
    path = config_path(directory)
    path.parent.mkdir(parents=True, exist_ok=True)
    ignore_self(path.parent)
    path.write_text(json.dumps({"url": url.rstrip("/"), "token": token}, indent=2))
    path.chmod(0o600)
    return path


def login(
    url: str,
    email: str,
    password: str,
    timeout: float = 30.0,
    directory: Path | None = None,
) -> Path:
    """Exchange credentials for a token and remember the engine."""
    import httpx

    response = httpx.post(
        f"{url.rstrip('/')}{API}/login/access-token",
        data={"username": email, "password": password},
        timeout=timeout,
    )
    if response.status_code >= 400:
        raise ApiError(response.status_code, _detail(response))
    return write_config(url, str(response.json()["access_token"]), directory)


# ---------------------------------------------------------------------------
# sync
# ---------------------------------------------------------------------------


class SyncReport:
    """What one flow's sync did, for the CLI to print."""

    def __init__(self, flow: str) -> None:
        self.flow = flow
        self.created = False
        self.changed: list[str] = []
        self.published = False

    @property
    def unchanged(self) -> bool:
        return not self.created and not self.changed

    def __repr__(self) -> str:
        return f"SyncReport({self.flow!r}, changed={self.changed})"


def origin_of(repo: str) -> dict[str, Any]:
    """Where this flow came from, so a run can name the code that produced it."""
    commit, dirty = "", False
    if repo:
        head = _git(repo, "rev-parse", "HEAD")
        if head is not None:
            commit = head
            dirty = bool(_git(repo, "status", "--porcelain"))
    return {"kind": "python", "repo": repo, "commit": commit, "dirty": dirty}


def _git(repo: str, *args: str) -> str | None:
    try:
        result = subprocess.run(
            ["git", "-C", repo, *args], capture_output=True, text=True, timeout=10
        )
    except (OSError, subprocess.SubprocessError):
        return None
    return result.stdout.strip() if result.returncode == 0 else None


def repo_root(start: str | Path) -> str:
    """The git repository a path sits in, or its directory if there is none."""
    directory = Path(start).resolve()
    if directory.is_file():
        directory = directory.parent
    found = _git(str(directory), "rev-parse", "--show-toplevel")
    return found or str(directory)


def sync(
    flows: Iterable[Flow],
    client: Client,
    *,
    repo: str = "",
    origin: dict[str, Any] | None = None,
    publish: bool = True,
    force: bool = False,
) -> list[SyncReport]:
    """Put declared flows into the store, refusing to overwrite canvas work.

    The order matters: the document first, so the nodes it names exist, then
    each node's generated body, then one publish. Both writes no-op on
    identical content, so a second sync with nothing changed commits nothing.
    """
    stamp = origin if origin is not None else origin_of(repo)
    reports = []
    for target in flows:
        reports.append(_sync_one(target, client, stamp, publish=publish, force=force))
    # Always, even when nothing here changed. A worker holds the imported
    # package in `sys.modules` for as long as it lives, so an edit to the
    # caller's own code is invisible until the process is retired — and that
    # edit is invisible to this function too, since it changes no shim and no
    # document. Refusing to refresh on a "no-op" sync would skip exactly the
    # case the refresh exists for.
    client.refresh_modules()
    return reports


def _sync_one(
    target: Flow,
    client: Client,
    origin: dict[str, Any],
    *,
    publish: bool,
    force: bool,
) -> SyncReport:
    report = SyncReport(target.name)
    stored = client.get_flow(target.name)
    version = 1
    if stored is None:
        report.created = True
    else:
        definition = stored.get("definition") or {}
        version = int(definition.get("version") or 1)
        if not force:
            _refuse_on_drift(client, target, definition)

    saved = client.put_flow(target.document(origin) | {"version": version})
    stored_version = int((saved.get("definition") or {}).get("version") or version)
    if report.created or stored_version != version:
        # The store bumps the version only when the content actually changed,
        # so this is its own no-op detection rather than a second guess at it.
        report.changed.append("flow")
    version = stored_version

    for node_id, code in target.shims().items():
        if not report.created and client.get_source(target.name, node_id) == code:
            continue
        client.put_source(target.name, node_id, code)
        report.changed.append(node_id)

    if publish and (client.get_flow(target.name) or {}).get("has_draft"):
        client.publish(target.name, version)
        report.published = True
    return report


def _refuse_on_drift(client: Client, target: Flow, definition: dict[str, Any]) -> None:
    """Stop before overwriting work that was done somewhere else.

    Two ways a stored flow is not ours to replace: it was drawn on the canvas
    and has no origin at all, or one of its node bodies no longer carries the
    line saying it was generated — which means somebody edited the code there.
    """
    if not definition.get("origin"):
        raise SyncError(
            f"flow '{target.name}' was not created by sync, so replacing it would "
            "discard whoever drew it. Rename yours, or pass --force."
        )
    for stored_node in definition.get("nodes") or []:
        node_id = str(stored_node.get("id"))
        try:
            code = client.get_source(target.name, node_id)
        except ApiError:
            continue
        if code and not code.startswith(MARKER):
            raise SyncError(
                f"node '{target.name}.{node_id}' was edited on the canvas, and "
                "syncing would throw that edit away. Copy it out, or pass --force."
            )