"""Shared setup for the datavo store-operations examples.

Every example script builds its client with :func:`make_client`, so the flow
code stays focused on the operation it demonstrates rather than on auth.

Pick the server with the ``DATAVO_SERVER`` environment variable; if unset, the
active profile in ``~/.datavo/config.json`` is used (falling back to a local
dev server). Authentication is chosen from the server's ``/config``:

* ``simple`` auth (the local dev default) needs no token at all.
* ``entra`` / ``easyauth`` (dev, prod) authenticate through MSAL: the token is
  refreshed silently from ``~/.datavo/msal_cache.json`` when a cached refresh
  token is available, and falls back to an interactive browser / device-code
  login otherwise (run ``datavo login`` once to prime the cache).

Run an example against dev::

    DATAVO_SERVER=https://dev.datavo.io python examples/01_ingest_source.py

…or against a local ``uv run datavo-api`` (simple auth, no token)::

    DATAVO_SERVER=http://localhost:8000 python examples/01_ingest_source.py
"""

from __future__ import annotations

import io
import json
import os
import uuid
from pathlib import Path

from datavo_sdk import (
    DatasetResponse,
    SourceSpec,
    DatavoClient,
    KeyGroupSpec,
    ShardPlan,
    TarWriter,
    WorksetCreateRequest,
)


def resolve_server() -> str:
    """Server URL from ``DATAVO_SERVER``, else the active CLI profile, else local.

    Note the fallback is **localhost, not the SDK's production default**, and that is
    deliberate: every example creates and then deletes a store, so an unconfigured run
    must not quietly do that against production. Point at a real server explicitly.
    """
    env = (os.environ.get("DATAVO_SERVER") or "").strip()
    if env:
        return env
    config_path = Path.home() / ".datavo" / "config.json"
    if config_path.exists():
        try:
            payload = json.loads(config_path.read_text(encoding="utf-8"))
            active = (os.environ.get("DATAVO_PROFILE") or "").strip() or payload.get("active_profile")
            profile = (payload.get("profiles") or {}).get(active) or {}
            url = (profile.get("api_base_url") or "").strip()
            if url:
                return url
        except (json.JSONDecodeError, OSError):
            pass
    return "http://localhost:8000"


def make_client(server: str | None = None) -> DatavoClient:
    """Build a :class:`DatavoClient` for ``server`` with the right auth.

    ``simple``-auth servers get a token-less client; ``entra``/``easyauth``
    servers get the SDK's silent-refresh login (interactive fallback).
    """
    server = server or resolve_server()
    # ``/config`` is public — read it without triggering a login.
    auth_provider = DatavoClient(server_url=server, auto_login=False).get_public_config().auth_provider
    if auth_provider == "simple":
        return DatavoClient(server_url=server, auto_login=False)
    return DatavoClient(server_url=server)


def step(message: str) -> None:
    """Print a numbered-looking progress line so example output reads cleanly."""
    print(f"  → {message}", flush=True)


def unique(base: str) -> str:
    """A collision-safe store name for a demo run (``base`` + short suffix)."""
    return f"{base}_{uuid.uuid4().hex[:6]}"


def build_source_tar(samples: list[dict], *, root_key: str) -> bytes:
    """Build a WebDataset source shard: one member per ``{sample_key}.{field}``,
    behind the conventional ``__meta__.json`` header (optional for import — the
    server skips it — but written here so the tars are ordinary WebDataset shards;
    ``root_key`` is what that header advertises to other WebDataset readers).
    ``samples`` are ``{"__key__": id, key: value, ...}``.
    """
    buffer = io.BytesIO()
    meta = {
        "metadata_version": 1,
        "created_by": {"type": "example"},
        "root": root_key,
        "key_derivation": {},
        "tags": [],
    }
    with TarWriter(buffer, tar_meta=meta) as writer:
        for sample in samples:
            writer.write(sample)
    return buffer.getvalue()


def seed_source_store(
    client: DatavoClient,
    store: str,
    *,
    samples: list[dict],
    root_key: str = "audio.wav",
    default_tags: dict[str, str] | None = None,
):
    """Create ``store`` and import ``samples`` into it via a source import.

    This is the ingest flow spelled out in ``01_ingest_source.py``, wrapped up so
    the other examples can get a store full of samples in one call.
    """
    client.create_store(name=store)
    imp = client.create_source_import(
        sample_store=store,
        operation_id=f"{store}-src",
        default_tags=default_tags or {},
    )
    client.upload_source_import_part(imp.source_import_id, 0, build_source_tar(samples, root_key=root_key))
    client.commit_source_import(imp.source_import_id)
    return client.wait_for_source_import(imp.source_import_id, timeout_seconds=180)


def wait_for_split_set(
    client: DatavoClient,
    dataset: str,
    split_set: str,
    *,
    state: str = "ready",
    timeout_seconds: float = 600.0,
    poll_interval_seconds: float = 2.0,
):
    """Poll a split set until it reaches ``state`` (default the terminal ``ready``).

    ``create_dataset_split_set`` is worker-planned like ``create_dataset``: the POST
    returns a ``planning`` split set with no children yet, and the worker cuts the
    train/val/test child datasets asynchronously. The SDK ships waiters for imports,
    Adds, and datasets but not split sets, so the examples poll here. Raises on
    ``failed`` / timeout, mirroring ``client.wait_for_dataset``.
    """
    import time

    deadline = time.monotonic() + timeout_seconds
    while True:
        detail = client.get_dataset_split_set(dataset, split_set)
        if detail.lifecycle_state == state:
            return detail
        if detail.lifecycle_state == "failed" and state != "failed":
            raise RuntimeError(
                f"split set {split_set!r} on {dataset!r} failed: {detail.error_message or 'unknown error'}"
            )
        if time.monotonic() >= deadline:
            raise TimeoutError(
                f"timed out waiting for split set {split_set!r} on {dataset!r} to become {state}"
            )
        time.sleep(poll_interval_seconds)


def cleanup_split_sets(client: DatavoClient, dataset: str, *split_sets: str) -> None:
    """Best-effort teardown for split sets: archive then delete each one.

    A split set must be archived before it can be deleted, and deleting it cascades
    its child (train/val/test) datasets. Tear split sets down **before** their parent
    dataset (a parent with live split sets refuses deletion) and after any downstream
    store that was seeded from a split (that Add pins the split child). Failures are
    reported, not raised."""
    for split_set in split_sets:
        try:
            client.archive_dataset_split_set(dataset, split_set)
        except Exception:
            pass  # already archived, or never created
        try:
            client.delete_dataset_split_set(dataset, split_set)
            print(f"  ✓ cleaned up split set {split_set!r} (and its child splits)", flush=True)
        except Exception as exc:  # e.g. caller lacks permission to delete
            print(f"  note: left split set {split_set!r} archived (could not delete: {exc})", flush=True)


def build_dataset(
    client: DatavoClient,
    name: str,
    *,
    sources: list[SourceSpec],
    key_groups: list[KeyGroupSpec],
    shard_plan: ShardPlan,
    timeout_seconds: float = 600.0,
) -> DatasetResponse:
    """Build a dataset the canonical **workset-first** way.

    A workset carries the same composition a dataset does (sources union + key-groups), so
    the sanctioned way to build any dataset is: create a workset over the composition,
    (optionally explore it), then *convert* it into a permanent, reproducible dataset —
    ``create_workset(...)`` → ``convert_workset_to_dataset(...)``. Explore cheaply with a
    short-lived workset; keep the one you want.

    (Direct ``create_dataset(sources=…, key_groups=…)`` composition still works but is
    deprecated in favour of this path.) Returns the ready :class:`DatasetResponse`.
    """
    workset = client.create_workset(
        WorksetCreateRequest(sources=sources, key_groups=key_groups, shard_plan=shard_plan)
    )
    client.wait_for_workset(workset.workset_id, timeout_seconds=timeout_seconds)
    client.convert_workset_to_dataset(workset.workset_id, name, shard_plan=shard_plan)
    return client.wait_for_dataset(name, timeout_seconds=timeout_seconds)


def cleanup_datasets(client: DatavoClient, *names: str) -> None:
    """Best-effort teardown for datasets: archive then delete each one.

    A dataset holds a retention-tracked dependency on the store(s) it reads, so
    tear datasets down **before** the stores they were built over (call this
    ahead of :func:`cleanup`). Failures are reported, not raised."""
    for name in names:
        try:
            client.archive_dataset(name)
        except Exception:
            pass  # already archived, or never created
        try:
            client.delete_dataset(name)
            print(f"  ✓ cleaned up dataset {name!r}", flush=True)
        except Exception as exc:  # e.g. caller lacks permission to delete
            print(f"  note: left dataset {name!r} archived (could not delete: {exc})", flush=True)


def cleanup(client: DatavoClient, *stores: str) -> None:
    """Best-effort teardown: archive then delete each store (delete requires the
    store to be archived first). Failures are reported, not raised, so a demo
    always exits cleanly. Delete any datasets built over these stores first (see
    :func:`cleanup_datasets`) — a live dataset dependency blocks store deletion."""
    for name in stores:
        try:
            client.archive_store(name)
        except Exception:
            pass  # already archived, or never created
        try:
            client.delete_store(name)
            print(f"  ✓ cleaned up store {name!r}", flush=True)
        except Exception as exc:  # e.g. caller lacks permission to delete
            print(f"  note: left {name!r} archived (could not delete: {exc})", flush=True)