// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#include "RugnuxOutputFiles.h"

#include <filesystem>
#include <fstream>

#include <gemmi/ccp4.hpp>

#include "../writer/FileWriter.h"
#include "../writer/WriteModel.h"
#include "../writer/WriteRugnuxFiles.h"

namespace {

void WriteCcp4(const gemmi::Grid<float> &map, const std::string &path) {
    gemmi::Ccp4<float> ccp4;
    ccp4.grid = map;
    ccp4.update_ccp4_header(2);
    ccp4.write_ccp4_map(path);
}

// The files of one output pass at `prefix`, all but its _process.h5.
void WriteOutputs(const RugnuxOutputs &o, const std::string &prefix, const OutputProvenance &provenance,
                  size_t nthreads, Logger &logger) {
    if (o.merged)
        WriteMergedOutput(*o.merged, prefix, provenance, nthreads, /*hkl=*/true);
    if (o.twin)
        WriteMergedOutput(*o.twin, prefix, provenance, nthreads, /*hkl=*/false);
    if (o.model_prior)
        WriteMergedOutput(*o.model_prior, prefix, provenance, nthreads, /*hkl=*/false);
    if (o.maps)
        WriteModelMaps(*o.maps, prefix, logger);
    if (o.model_prior_maps)
        WriteModelMaps(*o.model_prior_maps, prefix + "_mr", logger);
    if (o.placed_model)
        WritePlacedModel(*o.placed_model->model, o.placed_model->cell, o.placed_model->space_group, prefix, logger);
    if (o.model_prior_placed_model)
        WritePlacedModel(*o.model_prior_placed_model->model, o.model_prior_placed_model->cell,
                         o.model_prior_placed_model->space_group, prefix + "_mr", logger);
    if (o.unmerged) {
        WriteMtz(*o.unmerged, prefix + "_unmerged.mtz");
        logger.Info("Unmerged observations written to {}_unmerged.mtz", prefix);
    }
    if (o.unmerged_partials) {
        WriteMtz(*o.unmerged_partials, prefix + "_unmerged_partials.mtz");
        logger.Info("Unmerged partials written to {}_unmerged_partials.mtz", prefix);
    }
    if (!o.detector_jpeg.empty()) {
        const std::string path = prefix + "_detector.jpg";
        std::ofstream f(path, std::ofstream::out | std::ofstream::trunc | std::ofstream::binary);
        if (f) {
            f.write(o.detector_jpeg.data(), static_cast<std::streamsize>(o.detector_jpeg.size()));
            logger.Info("Detector diagnostic image written to {}", path);
        } else
            logger.Warning("Cannot open {} for writing", path);
    }
    if (o.plot) {
        WritePerImagePlot(prefix, *o.plot);
        logger.Info("Per-image table written to {}_plot.txt", prefix);
    }
}

}  // namespace

ProcessFileSink::ProcessFileSink(std::string prefix) : prefix_(std::move(prefix)) {}

ProcessFileSink::~ProcessFileSink() = default;

void ProcessFileSink::Start(const StartMessage &start, const std::string &label) {
    writer_.reset();
    StartMessage s = start;
    s.file_prefix = prefix_ + label;
    writer_ = std::make_unique<FileWriter>(s, /*check_overwrite_at_start=*/true, /*trusted_path=*/true);
}

void ProcessFileSink::Write(const DataMessage &msg) {
    writer_->Write(msg);
}

void ProcessFileSink::End(const EndMessage &end) {
    writer_->WriteHDF5(end);
    writer_->Finalize();
    writer_.reset();
}

void WriteRugnuxOutputs(const ProcessResult &result, const std::string &prefix, const OutputProvenance &provenance,
                        size_t nthreads, Logger &logger) {
    if (result.outputs) {
        WriteOutputs(*result.outputs, prefix, provenance, nthreads, logger);
        if (result.outputs->process_start) {
            StartMessage start = *result.outputs->process_start;
            start.file_prefix = prefix;
            FileWriter writer(start, /*check_overwrite_at_start=*/true, /*trusted_path=*/true);
            writer.WriteHDF5(result.end_message);
            writer.Finalize();
        }
    }
    if (result.doubled_cell) {
        const ProcessResult::DoubledCell &h = *result.doubled_cell;
        if (h.other_outputs)
            WriteOutputs(*h.other_outputs, prefix + h.other_suffix, provenance, nthreads, logger);
        // The doubled cell's pass put its per-image analysis under a name of its own; where it is the answer,
        // the two _process.h5 files trade names, as the other files of the two passes are named.
        std::error_code ec;
        if (h.adopted && std::filesystem::exists(prefix + "_doubled_process.h5", ec)) {
            if (std::filesystem::exists(prefix + "_process.h5", ec))
                std::filesystem::rename(prefix + "_process.h5", prefix + h.other_suffix + "_process.h5", ec);
            std::filesystem::rename(prefix + "_doubled_process.h5", prefix + "_process.h5", ec);
        }
    }
}

void WriteMergedOutput(const MergedOutput &merged, const std::string &prefix, const OutputProvenance &provenance,
                       size_t nthreads, bool hkl) {
    // An MTZ, an mmCIF and a SHELX HKLF-4 .hkl - each has its uses downstream (MTZ for the CCP4 / phenix
    // reflection tools, mmCIF for deposition and as the self-describing native format, HKLF-4 for SHELXL
    // refinement and SHELXC / ANODE substructure solution). The .hkl holds the unmerged scaled fulls where the
    // merge provides them (rotation), the merged reflections otherwise (stills).
    const std::string file = prefix + merged.suffix;
    WriteMtz(merged.mtz, file + ".mtz");
    WriteMmcifReflections(merged.reflections, merged.cell, merged.experiment, merged.statistics, merged.error_model,
                          merged.twinning, provenance, file + ".cif", nthreads);
    if (!hkl)
        return;
    if (!merged.scaled_fulls.empty())
        WriteShelxHklReflections(merged.scaled_fulls, file + ".hkl", nthreads);
    else
        WriteShelxHklReflections(merged.reflections, merged.experiment, file + ".hkl", nthreads);
}

void WriteModelMaps(const ModelMaps &maps, const std::string &prefix, Logger &logger) {
    WriteMtz(maps.coefficients, prefix + "_maps.mtz");
    if (maps.map_2fofc)
        WriteCcp4(*maps.map_2fofc, prefix + "_2fofc.ccp4");
    if (maps.map_fofc)
        WriteCcp4(*maps.map_fofc, prefix + "_fofc.ccp4");
    if (maps.map_anom)
        WriteCcp4(*maps.map_anom, prefix + "_anom.ccp4");
    logger.Info("Model validation: wrote {}_maps.mtz{}", prefix,
                maps.map_2fofc ? fmt::format(", {0}_2fofc.ccp4, {0}_fofc.ccp4{1}", prefix,
                                             maps.map_anom ? ", " + prefix + "_anom.ccp4" : std::string())
                               : std::string());
}

void WritePerImagePlot(const std::string &prefix, const PerImagePlot &plot) {
    const std::string path = prefix + "_plot.txt";
    std::ofstream file(path, std::ofstream::out | std::ofstream::trunc);
    if (!file)
        throw JFJochException(JFJochExceptionCategory::FileWriteError, "Cannot open " + path + " for writing");
    file << PerImagePlotText(plot);
}

ReportDocument WriteResultReport(const std::string &output_prefix,
                                 const std::string &input_file,
                                 const DiffractionExperiment &experiment,
                                 const ProcessResult &result,
                                 Logger &logger,
                                 const RunProvenance &provenance) {
    ReportDocument doc = BuildReportDocument(output_prefix, input_file, experiment, result, provenance);
    if (output_prefix.empty())
        return doc;   // "compute the statistics, persist nothing"

    const std::string filename = output_prefix + "_report.txt";
    // The report is unconditional, so it must never be the reason a run fails: a run that produced a
    // good .mtz must survive an unwritable path or a full disk. Report the failure and carry on.
    try {
        std::ofstream file(filename);
        file.exceptions(std::ios::failbit | std::ios::badbit);
        file << RenderReportText(doc, provenance.developer);
    } catch (const std::exception &e) {
        logger.Warning("Could not write the results report {}: {}", filename, e.what());
    }
    return doc;
}