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

#include "MiniCBF.h"

#include <algorithm>
#include <cctype>
#include <cstring>
#include <fstream>
#include <limits>
#include <map>
#include <regex>
#include <stdexcept>

#include "../common/JFJochException.h"

namespace minicbf {

namespace {

// One capture group, first match, or nothing. The headers are a few kB, so a regex per field is
// cheap and keeps each rule next to the thing it reads.
std::optional<std::string> Match(const std::string &text, const char *pattern) {
    std::smatch m;
    const std::regex re(pattern);
    if (!std::regex_search(text, m, re) || m.size() < 2)
        return {};
    return m[1].str();
}

double Num(const std::string &text, const char *pattern, double fallback = 0.0) {
    const auto s = Match(text, pattern);
    return s.has_value() ? std::stod(*s) : fallback;
}

int64_t Int(const std::string &text, const char *pattern, int64_t fallback = 0) {
    const auto s = Match(text, pattern);
    return s.has_value() ? std::stoll(*s) : fallback;
}

// A goniometer angle, or nothing where the head has no such axis. Writers spell that -9999, and
// taking a sentinel for an angle would put the head somewhere it never was.
std::optional<double> Angle(const std::string &text, const char *pattern) {
    const auto s = Match(text, pattern);
    if (!s.has_value())
        return {};
    const double v = std::stod(*s);
    if (v < -9998.0)
        return {};
    return v;
}

// One "loop_" of the imgCIF template block these headers carry, as rows keyed by tag. Tags and values
// are both read as whitespace-separated tokens rather than by line, because a template packs several
// tags onto one line ("_axis.vector[1] _axis.vector[2] _axis.vector[3]") and the rows that follow are
// laid out to match the tags, not the lines.
std::vector<std::map<std::string, std::string>> ParseLoop(const std::string &text, const std::string &tag) {
    // The loop_ that introduces the tag, not the tag's own position: everything before it is another
    // loop's data.
    const size_t tag_at = text.find("\n" + tag);
    if (tag_at == std::string::npos)
        return {};
    const size_t loop_at = text.rfind("loop_", tag_at);
    if (loop_at == std::string::npos)
        return {};

    // Bare tokens, and the quoted ones a CIF value may be - a quoted value holding spaces would
    // otherwise be counted as several columns and shift every row after it.
    std::vector<std::string> tokens;
    for (size_t i = loop_at + 5; i < text.size();) {
        while ((i < text.size()) && std::isspace(static_cast<unsigned char>(text[i])))
            i++;
        if (i >= text.size())
            break;
        size_t end;
        if ((text[i] == '\'') || (text[i] == '"')) {
            end = text.find(text[i], i + 1);
            if (end == std::string::npos)
                break;
            tokens.push_back(text.substr(i + 1, end - i - 1));
            end++;
        } else {
            end = i;
            while ((end < text.size()) && !std::isspace(static_cast<unsigned char>(text[end])))
                end++;
            tokens.push_back(text.substr(i, end - i));
        }
        // A second loop_, or a tag belonging to another category, ends this one.
        if ((tokens.back() == "loop_")
                || (tokens.back().starts_with("_") && !tokens.back().starts_with(tag.substr(0, tag.find('.') + 1)))) {
            tokens.pop_back();
            break;
        }
        i = end;
    }

    std::vector<std::string> names;
    size_t first_value = 0;
    while ((first_value < tokens.size()) && tokens[first_value].starts_with("_"))
        names.push_back(tokens[first_value++]);
    if (names.empty())
        return {};

    std::vector<std::map<std::string, std::string>> rows;
    for (size_t i = first_value; i + names.size() <= tokens.size(); i += names.size()) {
        std::map<std::string, std::string> row;
        for (size_t j = 0; j < names.size(); j++)
            row[names[j]] = tokens[i + j];
        rows.push_back(std::move(row));
    }
    return rows;
}

std::string Field(const std::map<std::string, std::string> &row, const std::string &name) {
    const auto it = row.find(name);
    return (it == row.end()) ? std::string() : it->second;
}

// The imgCIF axis table: which axis turns or translates in which laboratory direction, and which two
// axes the image's columns and rows run along. Absent from most headers, which say nothing about any
// of this and are left exactly as they were read before.
//
// The element vectors are stated in the frame of the axis they depend on. Between them and the
// detector's own rotation every header seen has translations only, so they describe the image in the
// unswung detector frame - which is where the image orientation belongs, with the arm applied on top.
// Following Hammersley, Bernstein & Westbrook (2006) Int. Tables Cryst. G, 444-458
void ParseAxisTable(const std::string &text, Header &h) {
    const auto axes = ParseLoop(text, "_axis.id");
    if (axes.empty())
        return;

    std::map<std::string, std::array<double, 3>> vector_of;
    for (const auto &row: axes) {
        const std::string id = Field(row, "_axis.id");
        const std::string type = Field(row, "_axis.type");
        const std::string equipment = Field(row, "_axis.equipment");
        const std::string depends_on = Field(row, "_axis.depends_on");
        std::array<double, 3> v{};
        try {
            for (int i = 0; i < 3; i++)
                v[i] = std::stod(Field(row, "_axis.vector[" + std::to_string(i + 1) + "]"));
        } catch (const std::exception &) {
            continue;   // "." for a vector: the table states no direction for this axis
        }
        vector_of[id] = v;

        // The base spindle and the detector arm are the rotations that hang off nothing: everything
        // further in is carried by them. Naming neither, so a beamline is free to call them anything.
        if ((type == "rotation") && (depends_on == ".")) {
            if (equipment == "goniometer")
                h.spindle_axis = v;
            else if (equipment == "detector")
                h.detector_axis = v;
        }
    }

    // Which axis the fast index runs along, and which the slow, through the two tables that say so.
    std::map<std::string, std::string> axis_of_set;
    for (const auto &row: ParseLoop(text, "_array_structure_list_axis.axis_set_id"))
        axis_of_set[Field(row, "_array_structure_list_axis.axis_set_id")]
                = Field(row, "_array_structure_list_axis.axis_id");

    for (const auto &row: ParseLoop(text, "_array_structure_list.array_id")) {
        const std::string set = Field(row, "_array_structure_list.axis_set_id");
        const auto id = axis_of_set.contains(set) ? axis_of_set[set] : set;
        const auto it = vector_of.find(id);
        if (it == vector_of.end())
            continue;
        std::array<double, 3> v = it->second;
        if (Field(row, "_array_structure_list.direction") == "decreasing")
            for (double &c: v)
                c = -c;
        const std::string index = Field(row, "_array_structure_list.index");
        if (index == "1")
            h.fast_direction = v;
        else if (index == "2")
            h.slow_direction = v;
    }
}

int16_t ReadI16(const uint8_t *p) { int16_t v; std::memcpy(&v, p, 2); return v; }
int32_t ReadI32(const uint8_t *p) { int32_t v; std::memcpy(&v, p, 4); return v; }
int64_t ReadI64(const uint8_t *p) { int64_t v; std::memcpy(&v, p, 8); return v; }

}  // namespace

bool ScansPhi(const Header &h) {
    if (h.phi_increment_deg != 0.0 || h.omega_increment_deg != 0.0 || h.chi_increment_deg != 0.0)
        return h.phi_increment_deg != 0.0;
    std::string name = h.axis_name;
    std::transform(name.begin(), name.end(), name.begin(),
                   [](unsigned char c) { return std::tolower(c); });
    return name.starts_with("phi");
}

std::optional<size_t> FindBinarySection(const uint8_t *data, size_t size) {
    if (size < sizeof(BINARY_SEPARATOR))
        return {};
    const auto *end = data + size;
    const auto *hit = std::search(data, end, std::begin(BINARY_SEPARATOR), std::end(BINARY_SEPARATOR));
    if (hit == end)
        return {};
    return static_cast<size_t>(hit - data) + sizeof(BINARY_SEPARATOR);
}

Header ParseHeader(const char *data, size_t size) {
    const std::string t(data, size);
    Header h;

    h.detector       = Match(t, R"(#\s*Detector:\s*([^\r\n]+))").value_or("PILATUS");
    h.pixel_x_m      = Num(t, R"(#\s*Pixel_size\s+([\d.eE+-]+)\s*m)");
    h.pixel_y_m      = Num(t, R"(#\s*Pixel_size\s+[\d.eE+-]+\s*m\s*x\s*([\d.eE+-]+)\s*m)");
    // The unit is optional: one ALBA set writes "thickness 0.001000" with no " m" after it.
    h.thickness_m    = Num(t, R"(sensor,\s*thickness\s+([\d.eE+-]+))");
    h.distance_m     = Num(t, R"(#\s*Detector_distance\s+([\d.eE+-]+))");
    h.beam_x_px      = Num(t, R"(#\s*Beam_xy\s*\(\s*([\d.eE+-]+))");
    h.beam_y_px      = Num(t, R"(#\s*Beam_xy\s*\([^,]+,\s*([\d.eE+-]+))");
    h.wavelength_A   = Num(t, R"(#\s*Wavelength\s+([\d.eE+-]+))");
    h.start_angle_deg     = Num(t, R"(#\s*Start_angle\s+([\d.eE+-]+))");
    h.angle_increment_deg = Num(t, R"(#\s*Angle_increment\s+([\d.eE+-]+))");
    h.two_theta_deg  = Num(t, R"(#\s*Detector_2theta\s+([\d.eE+-]+))");
    // The whitespace after each name is what keeps "Chi_increment" out of "Chi".
    h.chi_deg        = Angle(t, R"(#\s*Chi\s+([\d.eE+-]+))");
    h.omega_deg      = Angle(t, R"(#\s*Omega\s+([\d.eE+-]+))");
    h.chi_increment_deg   = Angle(t, R"(#\s*Chi_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.phi_increment_deg   = Angle(t, R"(#\s*Phi_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.omega_increment_deg = Angle(t, R"(#\s*Omega_increment\s+([\d.eE+-]+))").value_or(0.0);
    h.exposure_s     = Num(t, R"(#\s*Exposure_time\s+([\d.eE+-]+))");
    h.period_s       = Num(t, R"(#\s*Exposure_period\s+([\d.eE+-]+))");
    h.count_cutoff   = Int(t, R"(#\s*Count_cutoff\s+(\d+))");
    h.axis_name      = Match(t, R"(#\s*Oscillation_axis\s+(\S+))").value_or("omega");
    // "+SLOW" / "+FAST" on that same line: which of the image's two directions the spindle runs
    // along. Some writers state that instead of an axis name, and it is the only thing a header with
    // no axis table says about the spindle's direction at all.
    h.spindle_along_slow = Match(t, R"(#\s*Oscillation_axis[^\r\n]*\+(SLOW|slow))").has_value();
    ParseAxisTable(t, h);

    // "# Silicon sensor, ..." / "# CdTe sensor, ...". The rest of the code compares the material
    // against "CdTe" (BraggIntegrationEngine), so an unnormalised "Silicon" would silently give a
    // CdTe sensor silicon's attenuation length.
    if (const auto m = Match(t, R"(#\s*(\w+)\s+sensor,)")) {
        std::string s = *m;
        std::transform(s.begin(), s.end(), s.begin(), [](unsigned char c) { return std::tolower(c); });
        h.material = (s == "cdte") ? "CdTe" : "Si";
    }

    h.nx    = Int(t, R"(X-Binary-Size-Fastest-Dimension:\s*(\d+))");
    h.ny    = Int(t, R"(X-Binary-Size-Second-Dimension:\s*(\d+))");
    h.nelem = Int(t, R"(X-Binary-Number-of-Elements:\s*(\d+))");

    const auto conv = Match(t, R"RE(conversions\s*=\s*"([^"]+)")RE").value_or("");
    h.byte_offset = conv.find("x-CBF_BYTE_OFFSET") != std::string::npos;

    if (h.period_s <= 0.0)
        h.period_s = h.exposure_s;

    return h;
}

// The x-CBF_BYTE_OFFSET scheme: a running value, each pixel stored as a delta in the smallest
// container that holds it, escaping to the next size with that container's most negative value.
// Following Bernstein & Hammersley (2006) Int. Tables Cryst. G, 37-43
void DecodeByteOffset(const uint8_t *data, size_t size, int32_t *out, size_t n_pixels) {
    int64_t value = 0;
    size_t pos = 0;
    size_t written = 0;

    while (written < n_pixels) {
        if (pos >= size)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  "miniCBF byte-offset stream ended after " + std::to_string(written)
                                  + " of " + std::to_string(n_pixels) + " pixels");

        int64_t delta = static_cast<int8_t>(data[pos]);
        pos += 1;

        if (delta == -128) {
            if (pos + 2 > size) break;
            delta = ReadI16(data + pos);
            pos += 2;
            if (delta == -32768) {
                if (pos + 4 > size) break;
                delta = ReadI32(data + pos);
                pos += 4;
                if (delta == INT32_MIN) {
                    if (pos + 8 > size) break;
                    delta = ReadI64(data + pos);
                    pos += 8;
                }
            }
        }

        value += delta;
        out[written++] = static_cast<int32_t>(value);
    }

    if (written != n_pixels)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "miniCBF byte-offset stream ended after " + std::to_string(written)
                              + " of " + std::to_string(n_pixels) + " pixels");
}

namespace {

std::vector<uint8_t> Slurp(const std::string &path, size_t max_bytes) {
    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Cannot open CBF file " + path);
    f.seekg(0, std::ios::end);
    const auto file_size = static_cast<size_t>(f.tellg());
    f.seekg(0, std::ios::beg);
    std::vector<uint8_t> buf(std::min(file_size, max_bytes));
    f.read(reinterpret_cast<char *>(buf.data()), static_cast<std::streamsize>(buf.size()));
    buf.resize(static_cast<size_t>(f.gcount()));
    return buf;
}

// Enough to reach the separator on any header seen in the wild (the longest measured is ~6.3 kB).
constexpr size_t HEADER_PROBE_BYTES = 256 * 1024;

}  // namespace

Header ReadHeader(const std::string &path) {
    const auto buf = Slurp(path, HEADER_PROBE_BYTES);
    const auto start = FindBinarySection(buf.data(), buf.size());
    if (!start.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No CBF binary section in " + path);
    return ParseHeader(reinterpret_cast<const char *>(buf.data()),
                       *start - sizeof(BINARY_SEPARATOR));
}

namespace {

// One read of the file, header parsed, dimensions checked. The caller supplies where the pixels go.
Header ReadCommon(const std::string &path, std::vector<uint8_t> &buf, size_t &binary_start) {
    buf = Slurp(path, std::numeric_limits<size_t>::max());
    const auto start = FindBinarySection(buf.data(), buf.size());
    if (!start.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "No CBF binary section in " + path);

    const Header h = ParseHeader(reinterpret_cast<const char *>(buf.data()),
                                 *start - sizeof(BINARY_SEPARATOR));
    if (!h.byte_offset)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Unsupported CBF compression in " + path + " (only x-CBF_BYTE_OFFSET)");
    if (h.nelem <= 0 || h.nx <= 0 || h.ny <= 0 || h.nelem != h.nx * h.ny)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "Inconsistent image dimensions in " + path);

    binary_start = *start;
    return h;
}

}  // namespace

Header Read(const std::string &path, std::vector<int32_t> &out) {
    std::vector<uint8_t> buf;
    size_t start = 0;
    const Header h = ReadCommon(path, buf, start);
    out.resize(static_cast<size_t>(h.nelem));
    DecodeByteOffset(buf.data() + start, buf.size() - start, out.data(),
                     static_cast<size_t>(h.nelem));
    return h;
}

Header ReadInto(const std::string &path, int32_t *out, size_t capacity) {
    std::vector<uint8_t> buf;
    size_t start = 0;
    const Header h = ReadCommon(path, buf, start);
    if (static_cast<size_t>(h.nelem) > capacity)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              "CBF image does not fit the supplied buffer: " + path);
    DecodeByteOffset(buf.data() + start, buf.size() - start, out, static_cast<size_t>(h.nelem));
    return h;
}

}  // namespace minicbf