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

#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;
}

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

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+-]+))");
    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");

    // "# 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;
}

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