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

#include "SFRM.h"

#include <algorithm>
#include <cctype>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <optional>
#include <sstream>

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

namespace sfrm {

namespace {

constexpr size_t BLOCK = 512;
constexpr size_t LINE = 80;
constexpr size_t KEY = 8;              // "NROWS  :" - seven characters of name and the colon

JFJochException Invalid(const std::string &what) {
    return {JFJochExceptionCategory::InputParameterInvalid, what};
}

std::vector<uint8_t> ReadPrefix(const std::string &path, size_t bytes) {
    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw JFJochException(JFJochExceptionCategory::MockFileOpenError, "Cannot open " + path);
    std::vector<uint8_t> out(bytes);
    f.read(reinterpret_cast<char *>(out.data()), static_cast<std::streamsize>(bytes));
    out.resize(static_cast<size_t>(f.gcount()));
    return out;
}

std::vector<std::string> Words(const std::map<std::string, std::string> &kv, const std::string &key) {
    std::vector<std::string> out;
    const auto it = kv.find(key);
    if (it == kv.end())
        return out;
    std::istringstream in(it->second);
    for (std::string w; in >> w;)
        out.push_back(w);
    return out;
}

double Num(const std::map<std::string, std::string> &kv, const std::string &key, size_t i,
           double fallback = 0.0) {
    const auto w = Words(kv, key);
    if (i >= w.size())
        return fallback;
    try {
        return std::stod(w[i]);
    } catch (const std::logic_error &) {
        throw Invalid("Malformed number in SFRM header " + key + ": '" + w[i] + "'");
    }
}

int64_t Int(const std::map<std::string, std::string> &kv, const std::string &key, size_t i,
            int64_t fallback = 0) {
    return static_cast<int64_t>(Num(kv, key, i, static_cast<double>(fallback)));
}

// The 80-character lines of the header, as key -> value. A key that repeats (TITLE, CELL, MATRIX,
// DETPAR, ...) continues its value on the next line, so the pieces are joined.
std::map<std::string, std::string> Lines(const char *data, size_t size) {
    std::map<std::string, std::string> kv;
    for (size_t at = 0; at + LINE <= size; at += LINE) {
        if (data[at + KEY - 1] != ':')
            continue;           // the padding after the last line
        std::string key(data + at, KEY - 1);
        key.erase(std::find_if(key.rbegin(), key.rend(), [](unsigned char c) { return !std::isspace(c); }).base(),
                  key.end());
        const std::string value(data + at + KEY, LINE - KEY);
        auto &v = kv[key];
        v += (v.empty() ? "" : " ") + value;
    }
    return kv;
}

bool HasExtension(const std::filesystem::path &p) {
    std::string ext = p.extension().string();
    std::transform(ext.begin(), ext.end(), ext.begin(), [](unsigned char c) { return std::tolower(c); });
    return ext == ".sfrm";
}

bool IsSFRM(const std::string &path) {
    const auto head = ReadPrefix(path, KEY);
    return head.size() == KEY && std::memcmp(head.data(), "FORMAT :", KEY) == 0;
}

// "ins_zwin_01_0042.sfrm" -> {"ins_zwin_01_", 4}: the run is part of the prefix, so naming a frame
// of run 01 reads run 01 and not the frames of every run in the directory.
struct Template {
    std::string prefix;
    size_t digits = 0;

    bool Matches(const std::string &stem) const {
        if (stem.size() != prefix.size() + digits || stem.compare(0, prefix.size(), prefix) != 0)
            return false;
        return std::all_of(stem.begin() + static_cast<std::ptrdiff_t>(prefix.size()), stem.end(),
                           [](unsigned char c) { return std::isdigit(c); });
    }
};

std::optional<Template> TemplateOf(const std::filesystem::path &p) {
    const std::string stem = p.stem().string();
    size_t start = stem.size();
    while (start > 0 && std::isdigit(static_cast<unsigned char>(stem[start - 1])))
        start--;
    if (start == stem.size())
        return {};
    return Template{stem.substr(0, start), stem.size() - start};
}

size_t Pad(size_t bytes, size_t to) {
    return (bytes + to - 1) / to * to;
}

uint64_t ReadLE(const uint8_t *p, int64_t bytes) {
    uint64_t v = 0;
    for (int64_t i = bytes - 1; i >= 0; i--)
        v = (v << 8) | p[i];
    return v;
}

}  // namespace

Header ParseHeader(const char *data, size_t size) {
    const auto kv = Lines(data, size);
    Header h;
    h.format = Int(kv, "FORMAT", 0);
    h.header_blocks = Int(kv, "HDRBLKS", 0);
    h.nrows = Int(kv, "NROWS", 0);
    h.ncols = Int(kv, "NCOLS", 0);
    h.bytes_per_pixel = Int(kv, "NPIXELB", 0, 1);
    h.bytes_per_underflow = Int(kv, "NPIXELB", 1, 1);
    if (h.format == 100) {
        for (size_t i = 0; i < 3; i++)
            h.overflows[i] = Int(kv, "NOVERFL", i, 0);
    } else {
        h.overflows = {Int(kv, "NOVERFL", 0, 0), 0, 0};
    }
    h.baseline = Int(kv, "NEXP", 2, 0);
    for (size_t i = 0; i < 4; i++)
        h.angles[i] = Num(kv, "ANGLES", i);
    h.axis = Int(kv, "AXIS", 0);
    h.start_deg = Num(kv, "START", 0);
    h.increment_deg = Num(kv, "INCREME", 0);
    const auto distance = Words(kv, "DISTANC");
    h.distance_cm = Num(kv, "DISTANC", distance.empty() ? 0 : distance.size() - 1);
    for (size_t i = 0; i < Words(kv, "CENTER").size(); i++)
        h.center.push_back(Num(kv, "CENTER", i));
    const auto dettype = Words(kv, "DETTYPE");
    h.detector = dettype.empty() ? "Bruker CCD" : dettype[0];
    h.pixels_per_cm_512 = Num(kv, "DETTYPE", 1);
    h.grid_offset_cm = Num(kv, "DETTYPE", 2);
    h.wavelength_A = Num(kv, "WAVELEN", 0);
    h.exposure_s = Num(kv, "CUMULAT", 0);
    h.saturation = Num(kv, "CCDPARM", 4);
    for (size_t i = 0; i < 6; i++)
        h.detpar[i] = Num(kv, "DETPAR", i);
    h.monochromator_2theta_deg = Num(kv, "FILTER2", 0);
    h.raw = kv;

    if (h.header_blocks <= 0 || h.nrows <= 0 || h.ncols <= 0)
        throw Invalid("SFRM header states no HDRBLKS, NROWS or NCOLS");
    if (h.bytes_per_pixel != 1 && h.bytes_per_pixel != 2 && h.bytes_per_pixel != 4)
        throw Invalid("SFRM header states " + std::to_string(h.bytes_per_pixel) + " bytes per pixel");
    if (Int(kv, "WORDORD", 0) != 0 || Int(kv, "LONGORD", 0) != 0)
        throw Invalid("big-endian SFRM frames are not supported");
    if (h.format != 86 && h.format != 100)
        throw Invalid("SFRM FORMAT " + std::to_string(h.format) + " is not supported");
    return h;
}

Header ReadHeader(const std::string &path) {
    // Enough to reach HDRBLKS (line 3), then exactly the blocks it states.
    auto buf = ReadPrefix(path, 5 * BLOCK);
    const auto blocks = Int(Lines(reinterpret_cast<const char *>(buf.data()), buf.size()), "HDRBLKS", 0);
    if (blocks <= 0)
        throw Invalid(path + " states no HDRBLKS");
    buf = ReadPrefix(path, static_cast<size_t>(blocks) * BLOCK);
    return ParseHeader(reinterpret_cast<const char *>(buf.data()), buf.size());
}

Header Decode(const uint8_t *data, size_t size, int32_t *out, size_t capacity) {
    const size_t header_bytes = 5 * BLOCK;
    if (size < header_bytes)
        throw Invalid("SFRM file shorter than its header");
    const auto blocks = Int(Lines(reinterpret_cast<const char *>(data), header_bytes), "HDRBLKS", 0);
    if (blocks <= 0 || static_cast<size_t>(blocks) * BLOCK > size)
        throw Invalid("SFRM file shorter than its header");
    const Header h = ParseHeader(reinterpret_cast<const char *>(data), static_cast<size_t>(blocks) * BLOCK);

    const size_t npixel = static_cast<size_t>(h.nrows) * static_cast<size_t>(h.ncols);
    if (npixel > capacity)
        throw Invalid("SFRM image does not fit the buffer given for it");
    size_t at = static_cast<size_t>(h.header_blocks) * BLOCK;
    if (at + npixel * h.bytes_per_pixel > size)
        throw Invalid("SFRM file ends before its image does");
    for (size_t i = 0; i < npixel; i++)
        out[i] = static_cast<int32_t>(ReadLE(data + at + i * h.bytes_per_pixel, h.bytes_per_pixel));
    at += npixel * h.bytes_per_pixel;

    if (h.format == 86) {
        // One ASCII entry per overflowing pixel: its count in 9 characters, its offset in 7.
        const int64_t n = h.overflows[0];
        if (n > 0 && at + static_cast<size_t>(n) * 16 > size)
            throw Invalid("SFRM file ends inside its overflow table");
        for (int64_t k = 0; k < n; k++) {
            const std::string entry(reinterpret_cast<const char *>(data + at + 16 * k), 16);
            int64_t value, offset;
            try {
                value = std::stoll(entry.substr(0, 9));
                offset = std::stoll(entry.substr(9, 7));
            } catch (const std::logic_error &) {
                throw Invalid("malformed SFRM overflow entry '" + entry + "'");
            }
            if (offset < 0 || static_cast<size_t>(offset) >= npixel)
                throw Invalid("SFRM overflow entry points outside the image");
            out[offset] = static_cast<int32_t>(value);
        }
    } else {
        at = Pad(at, 16);
        // The three tables, and where the next unread entry of each is.
        const auto table = [&](int64_t n, int64_t bytes) {
            const size_t begin = at;
            if (n > 0) {
                at = Pad(at + static_cast<size_t>(n) * bytes, 16);
                if (begin + static_cast<size_t>(n) * bytes > size)
                    throw Invalid("SFRM file ends inside its overflow tables");
            }
            return begin;
        };
        const size_t under = table(h.overflows[0], h.bytes_per_underflow);
        const size_t over16 = table(h.overflows[1], 2);
        const size_t over32 = table(h.overflows[2], 4);
        int64_t n_under = 0, n16 = 0, n32 = 0;
        const auto next = [&](size_t begin, int64_t &used, int64_t total, int64_t bytes) {
            if (used >= total)
                throw Invalid("SFRM frame has more overflowing pixels than its tables hold");
            return ReadLE(data + begin + static_cast<size_t>(used++) * bytes, bytes);
        };
        const uint64_t max_pixel = (h.bytes_per_pixel == 4) ? UINT32_MAX : (1ULL << (8 * h.bytes_per_pixel)) - 1;
        for (size_t i = 0; i < npixel; i++) {
            uint64_t v = static_cast<uint32_t>(out[i]);
            if (v == 0 && h.overflows[0] > 0)
                v = next(under, n_under, h.overflows[0], h.bytes_per_underflow);
            else if (h.bytes_per_pixel < 4 && v == max_pixel) {
                if (h.bytes_per_pixel == 1)
                    v = next(over16, n16, h.overflows[1], 2);
                if (v == 0xFFFF) {
                    const auto w = static_cast<int32_t>(next(over32, n32, h.overflows[2], 4));
                    out[i] = static_cast<int32_t>(w - h.baseline);
                    continue;
                }
            }
            out[i] = static_cast<int32_t>(static_cast<int64_t>(v) - h.baseline);
        }
        return h;
    }

    for (size_t i = 0; i < npixel; i++)
        out[i] -= static_cast<int32_t>(h.baseline);
    return h;
}

Header ReadInto(const std::string &path, int32_t *out, size_t capacity, std::vector<uint8_t> &scratch) {
    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw JFJochException(JFJochExceptionCategory::MockFileOpenError, "Cannot open " + path);
    f.seekg(0, std::ios::end);
    scratch.resize(static_cast<size_t>(f.tellg()));
    f.seekg(0, std::ios::beg);
    f.read(reinterpret_cast<char *>(scratch.data()), static_cast<std::streamsize>(scratch.size()));
    try {
        return Decode(scratch.data(), static_cast<size_t>(f.gcount()), out, capacity);
    } catch (const JFJochException &e) {
        throw Invalid(path + ": " + e.what());
    }
}

Header Read(const std::string &path, std::vector<int32_t> &out) {
    const Header h = ReadHeader(path);
    out.resize(static_cast<size_t>(h.nrows) * static_cast<size_t>(h.ncols));
    std::vector<uint8_t> scratch;
    return ReadInto(path, out.data(), out.size(), scratch);
}

std::vector<std::string> CollectSweep(const std::string &path) {
    const std::filesystem::path p(path);
    std::error_code ec;
    const bool is_dir = std::filesystem::is_directory(p, ec);
    std::filesystem::path dir = is_dir ? p : p.parent_path();
    if (dir.empty())
        dir = ".";

    // Naming a frame selects its run. Naming a directory selects the run with the most frames, which
    // is the one a user pointing at a data directory means - a Bruker collection directory holds one
    // run per scan, and every run is its own sweep.
    std::optional<Template> want;
    if (!is_dir)
        want = TemplateOf(p);

    std::map<std::pair<std::string, size_t>, std::vector<std::string>> runs;
    for (const auto &e: std::filesystem::directory_iterator(dir, ec)) {
        if (!e.is_regular_file() || !HasExtension(e.path()))
            continue;
        const auto t = TemplateOf(e.path());
        if (!t.has_value() || (want.has_value() && !want->Matches(e.path().stem().string())))
            continue;
        runs[{t->prefix, t->digits}].push_back(e.path().string());
    }
    std::vector<std::string> out;
    for (auto &[key, files]: runs)
        if (files.size() > out.size())
            out = std::move(files);
    std::sort(out.begin(), out.end());
    if (!out.empty() && !IsSFRM(out.front()))
        return {};
    return out;
}

bool CanRead(const std::string &path) {
    std::error_code ec;
    try {
        if (std::filesystem::is_directory(path, ec))
            return !CollectSweep(path).empty();
        return HasExtension(std::filesystem::path(path)) && IsSFRM(path);
    } catch (const JFJochException &) {
        return false;
    }
}

}  // namespace sfrm