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

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", "SOURCE_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"});
    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;
    }
    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;
}

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 {
        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]));
    }
    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