#include "SliceExport.hpp"

#include <fstream>

#include <vtkDataArray.h>
#include <vtkImageData.h>
#include <vtkNew.h>
#include <vtkPNGWriter.h>
#include <vtkPointData.h>

namespace geopro::app {

bool exportSliceImagePng(vtkImageData* colorImage, const std::string& path) {
    if (colorImage == nullptr || path.empty()) return false;
    vtkNew<vtkPNGWriter> writer;
    writer->SetFileName(path.c_str());
    writer->SetInputData(colorImage);  // 已上色 RGB 的切片 2D 图（非整窗截图）
    writer->Write();
    return writer->GetErrorCode() == 0;
}

bool exportSliceDat(vtkImageData* slice, const std::string& path) {
    if (slice == nullptr || path.empty()) return false;
    vtkDataArray* arr = slice->GetPointData() ? slice->GetPointData()->GetScalars() : nullptr;
    if (arr == nullptr) return false;
    int dims[3];
    slice->GetDimensions(dims);
    const int nx = dims[0], ny = dims[1];
    if (nx < 1 || ny < 1) return false;

    std::ofstream out(path);
    if (!out) return false;
    // 切片重采样为 2D（dims[2]=1）：写成行=j、列=i 的标量网格，每格取首分量。
    for (int j = 0; j < ny; ++j) {
        for (int i = 0; i < nx; ++i) {
            const vtkIdType id = static_cast<vtkIdType>(j) * nx + i;
            out << arr->GetComponent(id, 0) << (i + 1 < nx ? ' ' : '\n');
        }
    }
    return static_cast<bool>(out);
}

}  // namespace geopro::app