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

#include "WriteRugnuxFiles.h"

#include <chrono>
#include <stdexcept>
#include <filesystem>

#include <gemmi/mtz.hpp>

#include "WriteModel.h"
#include "../common/Logger.h"

void WriteMtz(const MtzTable &table, const std::string &filename) {
    const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(table.space_group_name);
    if (sg == nullptr)
        throw std::runtime_error(filename + ": unknown space group '" + table.space_group_name + "'");
    const size_t rows = table.columns.empty() ? 0 : table.columns[0].values.size();
    for (const auto &c : table.columns) {
        if (c.values.size() != rows)
            throw std::runtime_error(fmt::format("{}: column {} has {} rows, column {} {}", filename, c.label,
                                                 c.values.size(), table.columns[0].label, rows));
        if (c.dataset_id < 0 || c.dataset_id >= static_cast<int32_t>(table.datasets.size()))
            throw std::runtime_error(fmt::format("{}: column {} is in dataset {}, of {}", filename, c.label,
                                                 c.dataset_id, table.datasets.size()));
    }
    for (const auto &b : table.batches)
        if (b.ints.size() != 29 || b.floats.size() != 156
            || b.ints[20] < 0 || b.ints[20] >= static_cast<int32_t>(table.datasets.size()))
            throw std::runtime_error(fmt::format("{}: batch {} has no header of a dataset in the table", filename,
                                                 b.number));

    gemmi::Mtz mtz;
    mtz.title = table.title;
    mtz.history = table.history;
    mtz.spacegroup = sg;
    mtz.set_cell_for_all(table.unit_cell);
    mtz.sort_order = table.sort_order;
    for (size_t i = 0; i < table.datasets.size(); i++) {
        const MtzTable::Dataset &d = table.datasets[i];
        mtz.datasets.push_back({static_cast<int>(i), d.project_name, d.crystal_name, d.dataset_name, mtz.cell,
                                d.wavelength_A});
    }
    for (const auto &c : table.columns)
        mtz.add_column(c.label, c.type, c.dataset_id, -1, false);
    const size_t w = table.columns.size();
    mtz.nreflections = static_cast<int>(rows);
    mtz.data.resize(rows * w);
    for (size_t c = 0; c < w; c++)
        for (size_t r = 0; r < rows; r++)
            mtz.data[r * w + c] = table.columns[c].values[r];
    for (const auto &b : table.batches) {
        gemmi::Mtz::Batch &batch = mtz.batches.emplace_back();
        batch.number = b.number;
        batch.title = b.title;
        batch.ints.assign(b.ints.begin(), b.ints.end());
        batch.floats = b.floats;
        batch.axes = b.axes;
    }
    mtz.write_to_file(filename);
}

gemmi::Structure PlacedModelStructure(const PlacedModelMessage &model) {
    const size_t n = model.atom_name.size();
    for (const size_t size : {model.model_number.size(), model.chain.size(), model.residue_name.size(),
                              model.seq_num.size(), model.icode.size(), model.het_flag.size(), model.element.size(),
                              model.altloc.size(), model.charge.size(), model.x.size(), model.y.size(),
                              model.z.size(), model.b_iso.size(), model.occupancy.size()})
        if (size != n)
            throw std::runtime_error(fmt::format("placed model: a per-atom field has {} entries for {} atoms",
                                                 size, n));
    if (!model.aniso.empty() && model.aniso.size() != 6 * n)
        throw std::runtime_error(fmt::format("placed model: {} anisotropic values for {} atoms",
                                             model.aniso.size(), n));
    gemmi::Structure st;
    st.name = model.name;
    st.cell = model.unit_cell;
    st.spacegroup_hm = model.model_space_group_name;
    const bool has_aniso = !model.aniso.empty();
    gemmi::Model *m = nullptr;
    gemmi::Chain *ch = nullptr;
    gemmi::Residue *res = nullptr;
    for (size_t i = 0; i < model.atom_name.size(); i++) {
        if (m == nullptr || m->num != model.model_number[i]) {
            m = &st.models.emplace_back(model.model_number[i]);
            ch = nullptr;
        }
        if (ch == nullptr || ch->name != model.chain[i]) {
            ch = &m->chains.emplace_back(model.chain[i]);
            res = nullptr;
        }
        const gemmi::SeqId seqid(model.seq_num[i], static_cast<char>(model.icode[i]));
        if (res == nullptr || res->name != model.residue_name[i] || res->seqid != seqid
            || res->het_flag != static_cast<char>(model.het_flag[i])) {
            res = &ch->residues.emplace_back(gemmi::ResidueId{seqid, {}, model.residue_name[i]});
            res->het_flag = static_cast<char>(model.het_flag[i]);
        }
        gemmi::Atom &a = res->atoms.emplace_back();
        a.name = model.atom_name[i];
        a.element = gemmi::Element(model.element[i]);
        a.altloc = static_cast<char>(model.altloc[i]);
        a.charge = static_cast<signed char>(model.charge[i]);
        a.pos = gemmi::Position(model.x[i], model.y[i], model.z[i]);
        a.b_iso = model.b_iso[i];
        a.occ = model.occupancy[i];
        if (has_aniso)
            a.aniso = {model.aniso[6 * i], model.aniso[6 * i + 1], model.aniso[6 * i + 2],
                       model.aniso[6 * i + 3], model.aniso[6 * i + 4], model.aniso[6 * i + 5]};
    }
    st.setup_cell_images();
    return st;
}

namespace {

double WallClockSeconds() {
    return std::chrono::duration<double>(std::chrono::system_clock::now().time_since_epoch()).count();
}

// Rename tmp onto final, unless final is there already and may not be overwritten.
void RenameIntoPlace(const std::string &tmp, const std::string &final_name, bool overwrite, Logger &logger) {
    if (!overwrite && std::filesystem::exists(final_name)) {
        std::filesystem::remove(tmp);
        logger.Error("{} already exists and overwrite is off - not written", final_name);
        return;
    }
    std::filesystem::rename(tmp, final_name);
    logger.Info("{} written, renamed into place at wall clock {:.3f} s", final_name, WallClockSeconds());
}

}  // namespace

void WriteRugnuxResultFiles(const EndMessage &end, const std::string &file_prefix, bool overwrite,
                            Logger &logger) {
    // A file that cannot be written costs only itself: the master is already in place.
    try {
        if (end.merged) {
            WriteMtz(*end.merged, file_prefix + ".mtz.tmp");
            RenameIntoPlace(file_prefix + ".mtz.tmp", file_prefix + ".mtz", overwrite, logger);
        }
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}.mtz: {}", file_prefix, e.what());
    }
    try {
        if (end.map_coefficients) {
            WriteMtz(*end.map_coefficients, file_prefix + "_maps.mtz.tmp");
            RenameIntoPlace(file_prefix + "_maps.mtz.tmp", file_prefix + "_maps.mtz", overwrite, logger);
        }
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}_maps.mtz: {}", file_prefix, e.what());
    }
    try {
        if (end.placed_model) {
            const PlacedModelMessage &model = *end.placed_model;
            const gemmi::SpaceGroup *sg = gemmi::find_spacegroup_by_name(model.space_group_name);
            if (sg == nullptr)
                throw std::runtime_error("unknown space group '" + model.space_group_name + "'");
            // WritePlacedModel appends _model.cif and _model.pdb to the prefix it is given.
            const std::string tmp_prefix = file_prefix + ".tmp";
            WritePlacedModel(PlacedModelStructure(model), model.unit_cell, *sg, tmp_prefix, logger);
            for (const std::string suffix : {"_model.cif", "_model.pdb"})
                if (std::filesystem::exists(tmp_prefix + suffix))
                    RenameIntoPlace(tmp_prefix + suffix, file_prefix + suffix, overwrite, logger);
        }
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}_model.cif: {}", file_prefix, e.what());
    }
}