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

#include "MarCCD.h"

#include <tiffio.h>

#include <algorithm>
#include <cctype>
#include <cstring>
#include <filesystem>
#include <fstream>
#include <map>
#include <optional>
#include <tuple>

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

namespace marccd {

namespace {

// The instrument header, as marCCD writes it: 3072 bytes of little-endian int32 at fixed offsets.
// Only the fields below are read; the rest are detector housekeeping. Offsets are from the start of
// the block, which is where the "MMX" name sits.
constexpr size_t HEADER_BYTES = 3072;
constexpr size_t OFF_NAME = 4;             // char[16], "MMX"
constexpr size_t OFF_NFAST = 80;
constexpr size_t OFF_NSLOW = 84;
constexpr size_t OFF_DEPTH = 88;           // bytes per pixel
constexpr size_t OFF_SATURATED = 104;
// Goniostat block: distance and beam centre, then the driven circles as start/end pairs.
constexpr size_t OFF_XTAL_TO_DETECTOR = 640;   // micrometres
constexpr size_t OFF_BEAM_X = 644;             // 1/1000 pixel, along fast
constexpr size_t OFF_BEAM_Y = 648;             // 1/1000 pixel, along slow
constexpr size_t OFF_START_XTAL_TO_DETECTOR = 696;   // the same distance at the start of the sweep
constexpr size_t OFF_END_XTAL_TO_DETECTOR = 728;     // ... and at its end
constexpr size_t OFF_EXPOSURE_TIME = 656;      // milliseconds
constexpr size_t OFF_START_TWOTHETA = 668;     // millidegrees, and the seven that follow it
constexpr size_t OFF_END_TWOTHETA = 700;       // the same seven, at the end of the exposure
constexpr size_t OFF_ROTATION_RANGE = 736;     // millidegrees
constexpr size_t OFF_PIXELSIZE_X = 772;        // nanometres
constexpr size_t OFF_PIXELSIZE_Y = 776;
constexpr size_t OFF_SOURCE_WAVELENGTH = 908;  // 1e-5 angstrom

// The circles of the goniostat block, in the order it stores them. The scanned one is whichever
// pair of start/end angles differs, so this table is what turns that into a name.
constexpr const char *CIRCLE_NAMES[] = {"two_theta", "omega", "chi", "kappa",
                                        "phi", "delta", "gamma"};
constexpr size_t N_CIRCLES = sizeof(CIRCLE_NAMES) / sizeof(CIRCLE_NAMES[0]);

int32_t I32(const std::vector<uint8_t> &h, size_t off) {
    int32_t v;
    std::memcpy(&v, h.data() + off, sizeof(v));
    return v;   // the header declares little-endian byte order and no writer has ever used the other
}

// The first bytes of a file: the TIFF header, the instrument header, and nothing else.
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;
}

// Where the instrument header begins. The file states it in a private TIFF tag; every writer seen
// puts it at 1024 and says so, but the tag is what the format defines, so it wins where present.
uint32_t HeaderOffset(const std::vector<uint8_t> &prefix) {
    if (prefix.size() < 8)
        return DEFAULT_HEADER_OFFSET;
    const bool little = prefix[0] == 'I' && prefix[1] == 'I';
    const auto u16 = [&](size_t o) -> uint32_t {
        if (o + 2 > prefix.size()) return 0;
        return little ? (uint32_t(prefix[o]) | uint32_t(prefix[o + 1]) << 8)
                      : (uint32_t(prefix[o]) << 8 | uint32_t(prefix[o + 1]));
    };
    const auto u32 = [&](size_t o) -> uint32_t {
        if (o + 4 > prefix.size()) return 0;
        return little ? (uint32_t(prefix[o]) | uint32_t(prefix[o + 1]) << 8 |
                         uint32_t(prefix[o + 2]) << 16 | uint32_t(prefix[o + 3]) << 24)
                      : (uint32_t(prefix[o]) << 24 | uint32_t(prefix[o + 1]) << 16 |
                         uint32_t(prefix[o + 2]) << 8 | uint32_t(prefix[o + 3]));
    };
    const uint32_t ifd = u32(4);
    const uint32_t n = u16(ifd);
    for (uint32_t i = 0; i < n; i++) {
        const size_t e = ifd + 2 + static_cast<size_t>(i) * 12;
        if (e + 12 > prefix.size())
            break;
        if (u16(e) == TIFFTAG_MARCCD_HEADER_OFFSET)
            return u32(e + 8);
    }
    return DEFAULT_HEADER_OFFSET;
}

// The instrument header block, or an empty vector where the file does not carry one.
std::vector<uint8_t> InstrumentHeader(const std::string &path) {
    const uint32_t off = HeaderOffset(ReadPrefix(path, DEFAULT_HEADER_OFFSET));
    auto prefix = ReadPrefix(path, off + HEADER_BYTES);
    if (prefix.size() < off + HEADER_BYTES)
        return {};
    std::vector<uint8_t> h(prefix.begin() + off, prefix.begin() + off + HEADER_BYTES);
    if (std::memcmp(h.data() + OFF_NAME, "MMX", 4) != 0)
        return {};
    return h;
}

Header Parse(const std::vector<uint8_t> &h, const std::string &path) {
    Header out;
    out.nx = I32(h, OFF_NFAST);
    out.ny = I32(h, OFF_NSLOW);
    if (out.nx <= 0 || out.ny <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": marCCD header states no image size");
    if (I32(h, OFF_DEPTH) != 2)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": only 16-bit marCCD images are supported");

    out.pixel_x_m = I32(h, OFF_PIXELSIZE_X) * 1e-9;
    out.pixel_y_m = I32(h, OFF_PIXELSIZE_Y) * 1e-9;
    // xtal_to_detector, or the start/end pair where that field was left at zero - several writers
    // (the BESSY Rayonix MX-225 sets among them) fill only the pair, and taking the first field
    // alone refused three otherwise perfectly good sweeps with "states no detector distance".
    int32_t dist = I32(h, OFF_XTAL_TO_DETECTOR);
    if (dist == 0) dist = I32(h, OFF_START_XTAL_TO_DETECTOR);
    if (dist == 0) dist = I32(h, OFF_END_XTAL_TO_DETECTOR);
    out.distance_m = dist * 1e-6;
    // The beam centre is stored in thousandths of a pixel, x along the fast direction and y along
    // the slow one, which is the same order the internal frame counts columns and rows in.
    out.beam_x_px = I32(h, OFF_BEAM_X) * 1e-3;
    out.beam_y_px = I32(h, OFF_BEAM_Y) * 1e-3;
    out.wavelength_A = I32(h, OFF_SOURCE_WAVELENGTH) * 1e-5;
    out.two_theta_deg = I32(h, OFF_START_TWOTHETA) * 1e-3;
    out.exposure_s = I32(h, OFF_EXPOSURE_TIME) * 1e-3;
    out.saturated_value = I32(h, OFF_SATURATED);

    // Which circle moved. The header also carries a rotation_axis index, but the start/end pair
    // that differs is the same answer read off the angles themselves, and it stays right on a file
    // whose index field was never filled in.
    out.angle_increment_deg = I32(h, OFF_ROTATION_RANGE) * 1e-3;
    for (size_t i = 0; i < N_CIRCLES; i++) {
        const double start = I32(h, OFF_START_TWOTHETA + 4 * i) * 1e-3;
        const double end = I32(h, OFF_END_TWOTHETA + 4 * i) * 1e-3;
        if (start != end) {
            out.axis_name = CIRCLE_NAMES[i];
            out.start_angle_deg = start;
            if (out.angle_increment_deg == 0.0)
                out.angle_increment_deg = end - start;
            break;
        }
    }
    return out;
}

// The extensions a marCCD frame is written under. A frame whose extension is the frame number
// (D1.001) is named by the mar software itself and is just as common as a named one.
bool PlausibleExtension(const std::filesystem::path &p) {
    std::string ext = p.extension().string();
    if (ext.empty())
        return false;
    ext.erase(ext.begin());                       // drop the dot
    if (std::all_of(ext.begin(), ext.end(), [](unsigned char c) { return std::isdigit(c); }))
        return true;
    std::transform(ext.begin(), ext.end(), ext.begin(),
                   [](unsigned char c) { return std::tolower(c); });
    return ext == "mccd" || ext == "img" || ext == "tif" || ext == "tiff" || ext == "marccd";
}

// The sweep a file belongs to, as a template: everything before the LAST run of digits in the whole
// file name, the length of that run, and whatever follows it. That one rule covers both naming
// schemes in use - "xtal_1_00042.mccd" -> {"xtal_1_", 5, ".mccd"} and "D1.042" -> {"D1.", 3, ""} -
// because in both the frame number is the last number in the name.
struct Template {
    std::string prefix;
    size_t digits = 0;
    std::string suffix;

    bool operator<(const Template &o) const {
        return std::tie(prefix, digits, suffix) < std::tie(o.prefix, o.digits, o.suffix);
    }
    bool Matches(const std::string &name) const {
        if (name.size() != prefix.size() + digits + suffix.size())
            return false;
        if (name.compare(0, prefix.size(), prefix) != 0)
            return false;
        if (name.compare(prefix.size() + digits, suffix.size(), suffix) != 0)
            return false;
        for (size_t i = 0; i < digits; i++)
            if (!std::isdigit(static_cast<unsigned char>(name[prefix.size() + i])))
                return false;
        return true;
    }
};

std::optional<Template> TemplateOf(const std::string &name) {
    size_t end = name.size();
    while (end > 0 && !std::isdigit(static_cast<unsigned char>(name[end - 1])))
        end--;
    if (end == 0)
        return {};                                // no number anywhere: not part of a sweep
    size_t start = end;
    while (start > 0 && std::isdigit(static_cast<unsigned char>(name[start - 1])))
        start--;
    return Template{name.substr(0, start), end - start, name.substr(end)};
}

}  // namespace

std::vector<std::string> CollectSweep(const std::string &path) {
    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 sweep; naming a directory selects the sweep with the most frames
    // in it, which is the one a user pointing at a data directory means.
    std::optional<Template> want;
    if (!is_dir)
        want = TemplateOf(p.filename().string());

    std::map<Template, std::vector<std::string>> sweeps;
    for (const auto &e : std::filesystem::directory_iterator(dir, ec)) {
        if (!e.is_regular_file() || !PlausibleExtension(e.path()))
            continue;
        const std::string name = e.path().filename().string();
        const auto t = TemplateOf(name);
        if (!t.has_value())
            continue;
        if (want.has_value() && !want->Matches(name))
            continue;
        sweeps[*t].push_back(e.path().string());
    }

    std::vector<std::string> out;
    for (auto &[key, files] : sweeps)
        if (files.size() > out.size())
            out = std::move(files);

    // The frame number is zero-padded in every mar naming scheme seen, so within one template a
    // plain sort is the collection order.
    std::sort(out.begin(), out.end());
    // A directory of TIFFs from anything else matches the name rule just as well, so the chosen
    // sweep is confirmed against the one thing only a marCCD file has.
    if (!out.empty() && InstrumentHeader(out.front()).empty())
        return {};
    return out;
}

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

Header ReadHeader(const std::string &path) {
    const auto h = InstrumentHeader(path);
    if (h.empty())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + " carries no marCCD instrument header");
    return Parse(h, path);
}

Header ReadInto(const std::string &path, int32_t *out, size_t capacity, std::vector<uint8_t> &scratch) {
    const Header header = ReadHeader(path);
    const size_t npixel = static_cast<size_t>(header.nx) * static_cast<size_t>(header.ny);
    if (npixel > capacity)
        throw JFJochException(JFJochExceptionCategory::ArrayOutOfBounds,
                              path + ": image is larger than the buffer given for it");

    TIFF *tiff = TIFFOpen(path.c_str(), "r");
    if (tiff == nullptr)
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError, "Cannot open TIFF " + path);
    struct Closer {
        TIFF *t;
        ~Closer() { TIFFClose(t); }
    } closer{tiff};

    uint32_t width = 0, length = 0;
    uint16_t bits = 0, samples = 1;
    TIFFGetField(tiff, TIFFTAG_IMAGEWIDTH, &width);
    TIFFGetField(tiff, TIFFTAG_IMAGELENGTH, &length);
    TIFFGetField(tiff, TIFFTAG_BITSPERSAMPLE, &bits);
    TIFFGetFieldDefaulted(tiff, TIFFTAG_SAMPLESPERPIXEL, &samples);
    if (width != header.nx || length != header.ny)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": TIFF image size differs from the marCCD header");
    if (bits != 16 || samples != 1)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": only single-channel 16-bit marCCD images are supported");

    const size_t scanline_bytes = static_cast<size_t>(width) * sizeof(uint16_t);
    if (scanline_bytes != static_cast<size_t>(TIFFScanlineSize(tiff)))
        throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError,
                              path + ": TIFFScanlineSize mismatch");
    scratch.resize(scanline_bytes * length);
    for (uint32_t row = 0; row < length; row++) {
        if (TIFFReadScanline(tiff, scratch.data() + static_cast<size_t>(row) * scanline_bytes,
                             row, 0) < 0)
            throw JFJochException(JFJochExceptionCategory::TIFFGeneratorError,
                                  path + ": TIFFReadScanline error");
    }

    // The stored pixels are unsigned 16-bit; the rest of the code works in signed 32-bit, so they
    // are widened here and nothing downstream has to know what the file held.
    const auto *src = reinterpret_cast<const uint16_t *>(scratch.data());
    for (size_t i = 0; i < npixel; i++)
        out[i] = static_cast<int32_t>(src[i]);
    return header;
}

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

}  // namespace marccd