// 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 <fstream>

#include <gemmi/mtz.hpp>
#include <gemmi/read_cif.hpp>   // cif::read_string
#include <gemmi/mmcif.hpp>   // make_structure_from_block
#include <nlohmann/json.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));
    for (const size_t size : {model.subchain.size(), model.entity_id.size(), model.label_seq.size()})
        if (size != 0 && size != n)
            throw std::runtime_error(fmt::format("placed model: a per-atom label 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]);
            if (!model.subchain.empty())
                res->subchain = model.subchain[i];
            if (!model.entity_id.empty())
                res->entity_id = model.entity_id[i];
            if (!model.label_seq.empty())
                res->label_seq.value = model.label_seq[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]};
    }
    if (!model.metadata_mmcif.empty()) {
        // The input model's metadata, from the mmCIF categories it came as: parsed by gemmi, never copied.
        gemmi::cif::Document doc = gemmi::cif::read_string(model.metadata_mmcif);
        if (doc.blocks.empty())
            throw std::runtime_error("placed model: the metadata block is empty");
        gemmi::Structure meta = gemmi::make_structure_from_block(doc.blocks[0]);
        st.entities = std::move(meta.entities);
        st.connections = std::move(meta.connections);
        st.cispeps = std::move(meta.cispeps);
        st.mod_residues = std::move(meta.mod_residues);
        st.helices = std::move(meta.helices);
        st.sheets = std::move(meta.sheets);
        st.assemblies = std::move(meta.assemblies);
        st.ncs = std::move(meta.ncs);
        st.meta = std::move(meta.meta);
        st.info = std::move(meta.info);
        st.resolution = meta.resolution;
    }
    st.setup_cell_images();
    return st;
}

std::string DensityBlobsJson(const DensityBlobsMessage &b) {
    const size_t n = b.kind.size();
    for (const size_t size : {b.x.size(), b.y.size(), b.z.size(), b.peak_sigma.size(), b.volume.size(), b.score.size(),
                              b.nearest_atom.size(), b.nearest_atom_distance.size(), b.polar_contacts.size(),
                              b.metal_contacts.size(), b.anomalous_sigma.size()})
        if (size != n)
            throw std::runtime_error(fmt::format("blobs: a field has {} entries for {} blobs", size, n));
    nlohmann::json list = nlohmann::json::array();
    for (size_t i = 0; i < n; i++)
        list.push_back({{"kind", b.kind[i]},
                        {"centroid", {b.x[i], b.y[i], b.z[i]}},
                        {"peak_sigma", b.peak_sigma[i]},
                        {"volume_A3", b.volume[i]},
                        {"score_sigma_A3", b.score[i]},
                        {"nearest_atom", b.nearest_atom[i]},
                        {"nearest_atom_distance_A", b.nearest_atom_distance[i]},
                        {"polar_contacts", b.polar_contacts[i]},
                        {"metal_contacts", b.metal_contacts[i]},
                        {"anomalous_sigma", b.anomalous_sigma[i]}});
    return nlohmann::json{{"blobs", list}}.dump(1) + "\n";
}

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());
}

// <file_prefix><name>_model.cif (and .pdb) from `model`.
void WriteModelFiles(const PlacedModelMessage &model, const std::string &file_prefix, const std::string &name,
                     bool overwrite, Logger &logger) {
    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 + name + ".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 + name + suffix, overwrite, logger);
}

}  // 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)
            WriteModelFiles(*end.placed_model, file_prefix, "", overwrite, logger);
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}_model.cif: {}", file_prefix, e.what());
    }
    try {
        if (end.refined_model)
            WriteModelFiles(*end.refined_model, file_prefix, "_refined", overwrite, logger);
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}_refined_model.cif: {}", file_prefix, e.what());
    }
    try {
        if (end.refined_map_coefficients) {
            WriteMtz(*end.refined_map_coefficients, file_prefix + "_refined_maps.mtz.tmp");
            RenameIntoPlace(file_prefix + "_refined_maps.mtz.tmp", file_prefix + "_refined_maps.mtz", overwrite,
                            logger);
        }
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}_refined_maps.mtz: {}", file_prefix, e.what());
    }
    try {
        if (end.blobs) {
            const std::string text = DensityBlobsJson(*end.blobs);
            {
                std::ofstream f(file_prefix + "_blobs.json.tmp", std::ofstream::out | std::ofstream::trunc);
                f.exceptions(std::ios::failbit | std::ios::badbit);
                f << text;
            }
            RenameIntoPlace(file_prefix + "_blobs.json.tmp", file_prefix + "_blobs.json", overwrite, logger);
        }
    } catch (const std::exception &e) {
        logger.Error("Cannot write {}_blobs.json: {}", file_prefix, e.what());
    }
}