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

#include "SMV.h"

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

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

namespace smv {

namespace {

// The brace block is never long - 512 or 1024 bytes in everything seen - but HEADER_BYTES states
// the real length and is itself inside the block, so read a generous prefix and trust the value.
constexpr size_t PROBE_BYTES = 8192;

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::string Trim(std::string s) {
    const auto ws = [](unsigned char c) { return std::isspace(c) != 0; };
    while (!s.empty() && ws(s.front())) s.erase(s.begin());
    while (!s.empty() && ws(s.back())) s.pop_back();
    return s;
}

// "{ KEY=value; KEY=value; }" -> the pairs. Nothing here is a CIF or a JSON: a key is everything
// before the first '=', a value everything to the next ';', and the block ends at the closing brace.
std::optional<std::map<std::string, std::string>> ParseBlock(const std::vector<uint8_t> &buf) {
    if (buf.empty() || buf.front() != '{')
        return {};
    const std::string text(reinterpret_cast<const char *>(buf.data()), buf.size());
    const size_t end = text.find('}');
    if (end == std::string::npos)
        return {};
    std::map<std::string, std::string> kv;
    size_t at = 1;
    while (at < end) {
        const size_t semi = text.find(';', at);
        if (semi == std::string::npos || semi > end)
            break;
        const std::string item = text.substr(at, semi - at);
        const size_t eq = item.find('=');
        if (eq != std::string::npos) {
            std::string key = Trim(item.substr(0, eq));
            std::transform(key.begin(), key.end(), key.begin(),
                           [](unsigned char c) { return std::toupper(c); });
            kv[key] = Trim(item.substr(eq + 1));
        }
        at = semi + 1;
    }
    return kv;
}

double Num(const std::map<std::string, std::string> &kv, const std::string &key, double dflt = 0.0) {
    const auto it = kv.find(key);
    if (it == kv.end()) return dflt;
    try { return std::stod(it->second); } catch (const std::exception &) { return dflt; }
}

// The first key of `keys` the header actually carries. SMV accumulated synonyms over twenty years
// of writers and the spellings are not interchangeable between files, only between vendors.
double NumAny(const std::map<std::string, std::string> &kv,
              std::initializer_list<const char *> keys, double dflt = 0.0) {
    for (const char *k : keys) {
        const auto it = kv.find(k);
        if (it != kv.end()) {
            try { return std::stod(it->second); } catch (const std::exception &) {}
        }
    }
    return dflt;
}

// A value holding several numbers, as d*TREK writes its vectors and circle lists.
std::vector<double> Nums(const std::map<std::string, std::string> &kv, const std::string &key) {
    std::vector<double> out;
    const auto it = kv.find(key);
    if (it == kv.end()) return out;
    std::istringstream in(it->second);
    double v;
    while (in >> v) out.push_back(v);
    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);
    std::string w;
    while (in >> w) out.push_back(w);
    return out;
}

// d*TREK, as Rigaku's CrystalClear writes it for the Saturn and R-AXIS families. Read the way dxtbx's
// FormatSMVRigakuSaturn reads it: the image directions are <det>DETECTOR_VECTORS combined by
// <det>SPATIAL_DISTORTION_VECTORS, the point of normal incidence (pixels) and the pixel size (mm) come
// from <det>SPATIAL_DISTORTION_INFO, and the detector circles from <det>GONIO_*, where the distance is
// a translation along the normal applied before the rotations - so the beam position stays the PONI
// and the distance the normal distance whatever the arm does.
//
// The distortion vectors are not optional. dxtbx's variant for headers without DTREK_DATE_TIME
// leaves them out, and on a Saturn 944+ sweep that puts the spindle 90 degrees off: DIALS indexes
// 12% of the spots with it and 98% with the vectors applied, at the deposited cell.
void ParseDtrek(const std::map<std::string, std::string> &kv, const std::string &path, Header &h) {
    const auto detector_names = Words(kv, "DETECTOR_NAMES");
    const std::string det = detector_names.empty() ? "" : detector_names[0];
    const auto info = Nums(kv, det + "SPATIAL_DISTORTION_INFO");
    const auto vectors = Nums(kv, det + "DETECTOR_VECTORS");
    const auto names = Words(kv, det + "GONIO_NAMES");
    const auto units = Words(kv, det + "GONIO_UNITS");
    const auto values = Nums(kv, det + "GONIO_VALUES");
    const auto circle_vectors = Nums(kv, det + "GONIO_VECTORS");
    const auto rotation = Nums(kv, "ROTATION");      // start, end, increment, exposure, ...
    const auto spindle = Nums(kv, "ROTATION_VECTOR");
    if (info.size() < 4 || vectors.size() < 6 || rotation.size() < 4 || spindle.size() < 3
        || names.size() != values.size() || units.size() != values.size()
        || circle_vectors.size() != 3 * values.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": d*TREK header does not describe the detector and the rotation");

    const auto type = kv.find("DATA_TYPE");
    if (type != kv.end() && type->second.find("short") == std::string::npos)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": only 16-bit SMV images are supported (Data_type=" + type->second + ")");

    h.beam_x_px = info[0];
    h.beam_y_px = info[1];
    h.pixel_x_m = info[2] * 1e-3;
    h.pixel_y_m = info[3] * 1e-3;
    // Each image direction as a combination of the two detector vectors: fast = d0 f + d1 s,
    // slow = d2 f + d3 s. Absent, the detector vectors are the image directions.
    auto d = Nums(kv, det + "SPATIAL_DISTORTION_VECTORS");
    if (d.size() < 4)
        d = {1, 0, 0, 1};
    std::array<double, 3> fast{}, slow{};
    for (int i = 0; i < 3; i++) {
        fast[i] = d[0] * vectors[i] + d[1] * vectors[3 + i];
        slow[i] = d[2] * vectors[i] + d[3] * vectors[3 + i];
    }
    h.fast_direction = fast;
    h.slow_direction = slow;

    h.distance_m = 0;
    for (size_t i = 0; i < values.size(); i++) {
        const std::array<double, 3> axis{circle_vectors[3 * i], circle_vectors[3 * i + 1],
                                         circle_vectors[3 * i + 2]};
        if (units[i] == "deg")
            h.detector_circles.push_back({axis, values[i]});
        else if (names[i] == "Distance")
            h.distance_m = values[i] * 1e-3;
        else if (values[i] != 0.0)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  path + ": d*TREK detector translation " + names[i] + " is not supported");
    }

    h.start_angle_deg = rotation[0];
    h.angle_increment_deg = rotation[2];
    h.exposure_s = rotation[3];
    h.spindle_axis = std::array<double, 3>{spindle[0], spindle[1], spindle[2]};
    const auto axis_name = Words(kv, "ROTATION_AXIS_NAME");
    if (!axis_name.empty()) {
        h.axis_name = axis_name[0];
        std::transform(h.axis_name.begin(), h.axis_name.end(), h.axis_name.begin(),
                       [](unsigned char c) { return std::tolower(c); });
    }

    // SOURCE_WAVELENGTH leads with HOW MANY wavelengths follow ("1.000000 1.541780"), so its first
    // number is a count, not a wavelength. SCAN_WAVELENGTH is the one the scan used.
    const auto scan_wavelength = Nums(kv, "SCAN_WAVELENGTH");
    const auto source_wavelength = Nums(kv, "SOURCE_WAVELENGTH");
    h.wavelength_A = !scan_wavelength.empty() ? scan_wavelength[0]
                   : (source_wavelength.size() >= 2 ? source_wavelength[1] : 0.0);

    h.overflow_ratio = static_cast<int64_t>(Num(kv, "RAXIS_COMPRESSION_RATIO"));
    const auto id = kv.find(det + "DETECTOR_IDENTIFICATION");
    if (id != kv.end())
        h.detector = id->second;
}

std::optional<std::map<std::string, std::string>> HeaderBlock(const std::string &path) {
    return ParseBlock(ReadPrefix(path, PROBE_BYTES));
}

Header Parse(const std::map<std::string, std::string> &kv, const std::string &path) {
    Header h;
    h.raw = kv;
    h.nx = static_cast<int64_t>(Num(kv, "SIZE1"));
    h.ny = static_cast<int64_t>(Num(kv, "SIZE2"));
    if (h.nx <= 0 || h.ny <= 0)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": SMV header states no image size");
    h.data_offset = static_cast<size_t>(Num(kv, "HEADER_BYTES", 512));

    const auto order = kv.find("BYTE_ORDER");
    h.little_endian = (order == kv.end()) || order->second.find("little") != std::string::npos;

    const auto type = kv.find("TYPE");
    if (type != kv.end() && type->second.find("short") == std::string::npos)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": only 16-bit SMV images are supported (TYPE=" + type->second + ")");
    h.bytes_per_pixel = 2;

    // A single PIXEL_SIZE is the usual spelling; the split pair appears on a few writers.
    const double px = NumAny(kv, {"PIXEL_SIZE", "PIXEL_SIZE_X", "PIXELSIZE"});
    const double py = NumAny(kv, {"PIXEL_SIZE_Y", "PIXEL_SIZE", "PIXELSIZE"});
    h.pixel_x_m = px * 1e-3;                       // millimetres in the file
    h.pixel_y_m = (py > 0 ? py : px) * 1e-3;
    h.distance_m = NumAny(kv, {"DISTANCE", "DETECTOR_DISTANCE"}) * 1e-3;

    // Millimetres in the file, pixels everywhere in this program. BEAM_CENTRE is the British
    // spelling some writers use; ADSC's own header uses BEAM_CENTER.
    const double bx_mm = NumAny(kv, {"BEAM_CENTER_X", "BEAM_CENTRE_X", "BEAM_X"});
    const double by_mm = NumAny(kv, {"BEAM_CENTER_Y", "BEAM_CENTRE_Y", "BEAM_Y"});
    h.beam_x_px = h.pixel_x_m > 0 ? bx_mm * 1e-3 / h.pixel_x_m : 0.0;
    h.beam_y_px = h.pixel_y_m > 0 ? by_mm * 1e-3 / h.pixel_y_m : 0.0;

    h.wavelength_A = NumAny(kv, {"WAVELENGTH"});
    h.angle_increment_deg = NumAny(kv, {"OSC_RANGE", "OSCILLATION_RANGE"});
    // OSC_START is the angle of THIS image; PHI is where the circle stands, which is the same
    // thing on the single-axis goniometers that write this format, and is the only value some
    // writers give.
    h.start_angle_deg = NumAny(kv, {"OSC_START", "PHI", "START_PHI", "OMEGA"});
    h.two_theta_deg = NumAny(kv, {"TWOTHETA", "TWO_THETA", "DETECTOR_2THETA"});
    h.exposure_s = NumAny(kv, {"TIME", "EXPOSURE_TIME"});
    h.saturation = static_cast<int64_t>(NumAny(kv, {"CCD_IMAGE_SATURATION", "SATURATED_VALUE"}));
    const auto sn = kv.find("DETECTOR_SN");
    h.detector = sn != kv.end() ? ("SMV detector S/N " + sn->second) : "SMV";
    // Which circle the file says moved. OSC_AXIS names it where present; otherwise these are
    // single-axis collections and the name is the conventional one.
    const auto axis = kv.find("OSC_AXIS");
    if (axis != kv.end() && !axis->second.empty()) {
        std::string a = axis->second;
        std::transform(a.begin(), a.end(), a.begin(), [](unsigned char c) { return std::tolower(c); });
        h.axis_name = a;
    }
    if (kv.count("DETECTOR_NAMES"))
        ParseDtrek(kv, path, h);
    return h;
}

// Extensions an SMV frame is written under. The content test below is what actually decides - .img
// is also used by miniCBF and by marCCD - so this is only a cheap pre-filter.
bool PlausibleExtension(const std::filesystem::path &p) {
    std::string ext = p.extension().string();
    if (ext.empty())
        return false;
    ext.erase(ext.begin());
    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 == "img" || ext == "smv" || ext == "osc";
}

// The same sweep rule as the marCCD reader: the last run of digits in the whole file name, so
// "xtal_1_00042.img" and "xtal.042" are both handled by one rule.
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 {};
    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 = ".";

    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);
    std::sort(out.begin(), out.end());
    // The name rule matches plenty that is not SMV, so the chosen sweep is confirmed against the
    // one thing only an SMV file has: a brace block declaring an image size.
    if (!out.empty()) {
        const auto kv = HeaderBlock(out.front());
        if (!kv.has_value() || !kv->count("SIZE1"))
            return {};
    }
    return out;
}

// Asked of every input before anything is read, including inputs that are not SMV at all, so it
// answers a question and never fails: a file it cannot open or make sense of is simply not one of
// ours. The directory arm has to be inside the try as much as the single-file one - it opens the
// first plausible file in the folder, which can be unreadable, and the throw came out of a call site
// that does not catch it, ending the run with no message at all.
bool CanRead(const std::string &path) {
    std::error_code ec;
    try {
        if (std::filesystem::is_directory(path, ec))
            return !CollectSweep(path).empty();
        if (!PlausibleExtension(std::filesystem::path(path)))
            return false;
        const auto kv = HeaderBlock(path);
        return kv.has_value() && kv->count("SIZE1") && kv->count("SIZE2");
    } catch (const JFJochException &) {
        return false;
    }
}

Header ReadHeader(const std::string &path) {
    const auto kv = HeaderBlock(path);
    if (!kv.has_value())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + " carries no SMV header block");
    return Parse(*kv, path);
}

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

    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw JFJochException(JFJochExceptionCategory::MockFileOpenError, "Cannot open " + path);
    f.seekg(static_cast<std::streamoff>(h.data_offset), std::ios::beg);
    scratch.resize(npixel * sizeof(uint16_t));
    f.read(reinterpret_cast<char *>(scratch.data()), static_cast<std::streamsize>(scratch.size()));
    if (static_cast<size_t>(f.gcount()) != scratch.size())
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              path + ": file ends before the image does");

    // Stored unsigned 16-bit, handed out as the signed 32-bit the rest of the code reads.
    const uint8_t *src = scratch.data();
    if (h.little_endian) {
        for (size_t i = 0; i < npixel; i++)
            out[i] = static_cast<int32_t>(static_cast<uint16_t>(src[2 * i] | (src[2 * i + 1] << 8)));
    } else {
        for (size_t i = 0; i < npixel; i++)
            out[i] = static_cast<int32_t>(static_cast<uint16_t>((src[2 * i] << 8) | src[2 * i + 1]));
    }
    if (h.overflow_ratio > 0) {
        for (size_t i = 0; i < npixel; i++)
            if (out[i] > 32767)
                out[i] = static_cast<int32_t>((out[i] - 32768) * h.overflow_ratio);
    }
    return h;
}

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

}  // namespace smv