import invariant from "tiny-invariant";

import { createClient } from "../client";
import type { ClientFn } from "../types/core";
import type { DatasetExamples, DatasetSelector } from "../types/datasets";
import { getDatasetInfoByName } from "./getDatasetInfoByName";

export type GetDatasetExamplesParams = ClientFn & {
  /** Dataset selector (ID, name, or version ID) */
  dataset: DatasetSelector;
};

/**
 * Get examples from a dataset
 * @param dataset - Dataset selector (ID, name, version ID, or splits)
 * @returns Dataset examples
 */
export async function getDatasetExamples({
  client: _client,
  dataset: datasetSelector,
}: GetDatasetExamplesParams): Promise<DatasetExamples> {
  const client = _client || createClient();

  let datasetId: string;

  if ("datasetName" in datasetSelector) {
    const datasetInfo = await getDatasetInfoByName({
      client,
      datasetName: datasetSelector.datasetName,
    });
    datasetId = datasetInfo.id;
  } else {
    datasetId = datasetSelector.datasetId;
  }

  const { versionId, splits } = datasetSelector;

  const response = await client.GET("/v1/datasets/{id}/examples", {
    params: {
      path: {
        id: datasetId,
      },
      query: {
        ...(versionId ? { version_id: versionId } : {}),
        ...(splits ? { split: splits } : {}),
      },
    },
  });

  invariant(response.data?.data, "Failed to get dataset examples");
  const examplesData = response.data.data;
  return {
    versionId: examplesData.version_id,
    examples: examplesData.examples.map(
      ({ node_id, updated_at, ...example }) => ({
        ...example,
        nodeId: node_id,
        updatedAt: new Date(updated_at),
      })
    ),
  };
}
