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

#include "RugnuxModelFile.h"

#include <fstream>
#include <iterator>
#include <stdexcept>

#include <zlib.h>
#include <nlohmann/json.hpp>
#include <gemmi/mmread.hpp>

namespace {

// The bytes of a gzip stream inflated, every member of it.
std::string Gunzip(const std::string &bytes) {
    z_stream s{};
    if (inflateInit2(&s, 15 + 32) != Z_OK)
        throw std::runtime_error("cannot start zlib");
    std::string out;
    char buf[1 << 16];
    s.next_in = reinterpret_cast<Bytef *>(const_cast<char *>(bytes.data()));
    s.avail_in = static_cast<uInt>(bytes.size());
    int ret = Z_OK;
    while (ret != Z_STREAM_END || s.avail_in > 0) {
        if (ret == Z_STREAM_END)
            inflateReset(&s);
        s.next_out = reinterpret_cast<Bytef *>(buf);
        s.avail_out = sizeof buf;
        ret = inflate(&s, Z_NO_FLUSH);
        if (ret != Z_OK && ret != Z_STREAM_END) {
            inflateEnd(&s);
            throw std::runtime_error("the gzip stream is damaged");
        }
        out.append(buf, sizeof buf - s.avail_out);
        if (ret == Z_OK && s.avail_in == 0 && s.avail_out != 0) {
            inflateEnd(&s);
            throw std::runtime_error("the gzip stream is cut short");
        }
    }
    inflateEnd(&s);
    return out;
}

std::string ReadFile(const std::string &path) {
    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw std::runtime_error("cannot open " + path);
    return {std::istreambuf_iterator<char>(f), std::istreambuf_iterator<char>()};
}

}  // namespace

RugnuxModel ParseRugnuxModel(const std::string &bytes, const std::string &name, const std::string &pae_json) {
    RugnuxModel m;
    m.name = name;
    std::string text = bytes.size() >= 2 && bytes[0] == '\x1f' && bytes[1] == '\x8b' ? Gunzip(bytes) : bytes;
    gemmi::cif::Document doc;
    m.structure = gemmi::read_structure_from_memory(text.data(), text.size(), name, gemmi::CoorFormat::Detect, &doc);
    if (m.structure.models.empty())
        throw std::runtime_error(name + " holds no atoms");
    m.modelcif = m.structure.input_format == gemmi::CoorFormat::Mmcif && !doc.blocks.empty()
                 && doc.blocks[0].has_mmcif_category("_ma_qa_metric");
    if (!pae_json.empty()) {
        try {
            const nlohmann::json j = nlohmann::json::parse(pae_json);
            const nlohmann::json &pae = j.is_array() ? j.at(0).at("predicted_aligned_error")
                                                     : j.at("predicted_aligned_error");
            std::vector<std::vector<double>> rows;
            for (const auto &row : pae)
                rows.push_back(row.get<std::vector<double>>());
            m.pae = std::move(rows);
        } catch (const std::exception &) {
        }
    }
    return m;
}

RugnuxModel ReadRugnuxModelFile(const std::string &path) {
    const std::string pae_path = PaePathFor(path);
    std::string pae_json;
    if (!pae_path.empty() && std::ifstream(pae_path))
        pae_json = ReadFile(pae_path);
    return ParseRugnuxModel(ReadFile(path), path, pae_json);
}

std::string PaePathFor(const std::string &model_path) {
    const size_t slash = model_path.find_last_of("/\\");
    const std::string dir = slash == std::string::npos ? std::string() : model_path.substr(0, slash + 1);
    const std::string name = model_path.substr(dir.size());
    const size_t at = name.rfind("-model_v");
    if (at == std::string::npos)
        return {};
    const size_t dot = name.find('.', at);
    const std::string version = name.substr(at + 8, dot == std::string::npos ? std::string::npos : dot - at - 8);
    return dir + name.substr(0, at) + "-predicted_aligned_error_v" + version + ".json";
}