import vtkImageData from "@kitware/vtk.js/Common/DataModel/ImageData";

export const extractAxialSliceAsImage = (
  volumeData: vtkImageData,
  sliceIndex: number
): string => {
  const dimensions = volumeData.getDimensions();

  if (sliceIndex < 0 || sliceIndex >= dimensions[2]) {
    throw new Error("Invalid slice index");
  }

  const scalars = volumeData.getPointData().getScalars();
  const slice = scalars
    .getData()
    .slice(
      sliceIndex * dimensions[0] * dimensions[1],
      (sliceIndex + 1) * dimensions[0] * dimensions[1]
    );

  const canvas = document.createElement("canvas");
  canvas.width = dimensions[0];
  canvas.height = dimensions[1];
  const context = canvas.getContext("2d")!;
  const imageDataObj = context.createImageData(dimensions[0], dimensions[1]);

  for (let i = 0; i < slice.length; i++) {
    const value = slice[i];
    imageDataObj.data[i * 4] = value;
    imageDataObj.data[i * 4 + 1] = value;
    imageDataObj.data[i * 4 + 2] = value;
    imageDataObj.data[i * 4 + 3] = 255;
  }

  context.putImageData(imageDataObj, 0, 0);

  return canvas.toDataURL();
};