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

#include "XTC.h"

#include <algorithm>
#include <cstring>
#include <filesystem>
#include <fstream>

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

namespace {

constexpr size_t DGRAM_HEADER = 40;   // clock time 8, timestamp 8, env 4, Xtc header 20
constexpr size_t XTC_HEADER = 20;
constexpr size_t FRAME_HEADER = 16;   // Camera::FrameV1: width, height, depth, offset
constexpr uint32_t ID_XTC = 1;
constexpr uint32_t ID_FRAME = 2;
constexpr uint32_t ID_EBEAM = 15;
constexpr uint32_t ID_RAYONIX_CONFIG = 73;
constexpr uint32_t TRANSITION_L1ACCEPT = 12;
constexpr uint32_t TRANSITION_LAST = 12;
constexpr size_t EBEAM_V7_SIZE = 4 + 20 * 8;
constexpr size_t EBEAM_V7_PHOTON_ENERGY = 4 + 17 * 8;

std::ifstream Open(const std::string &path) {
    std::ifstream f(path, std::ios::binary);
    if (!f)
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Cannot open " + path);
    return f;
}

bool ReadAt(std::ifstream &f, uint64_t pos, void *dst, size_t n) {
    f.clear();
    f.seekg(static_cast<std::streamoff>(pos));
    f.read(static_cast<char *>(dst), static_cast<std::streamsize>(n));
    return f.gcount() == static_cast<std::streamsize>(n);
}

uint32_t Word(const uint8_t *p, size_t i) {
    uint32_t v;
    std::memcpy(&v, p + 4 * i, 4);
    return v;
}

// The transition a datagram header names, from the top bits of its timestamp's first word.
uint32_t Transition(const uint8_t *dgram) {
    return (Word(dgram, 2) >> 24) & 0xf;
}

// A first datagram is a Configure with a container as its outermost Xtc, whose extent covers at
// least its own header.
bool PlausibleFirstDatagram(const uint8_t *h) {
    return Transition(h) <= TRANSITION_LAST && (Word(h, 8) & 0xffff) == ID_XTC && Word(h, 9) >= XTC_HEADER;
}

// Walks the Xtcs in [begin, end) of the file, adding what it knows to index: images and the beam only
// from an event (an L1Accept), the configuration from any transition.
void Walk(std::ifstream &f, const std::string &file, uint64_t begin, uint64_t end, bool event, xtc::Index &index) {
    uint64_t pos = begin;
    while (pos + XTC_HEADER <= end) {
        uint8_t h[XTC_HEADER];
        if (!ReadAt(f, pos, h, XTC_HEADER))
            return;
        const uint32_t damage = Word(h, 0);
        const uint64_t source = (static_cast<uint64_t>(Word(h, 1)) << 32) | Word(h, 2);
        const uint32_t type = Word(h, 3) & 0xffff;
        const uint32_t version = (Word(h, 3) >> 16) & 0x7fff;
        const bool compressed = (Word(h, 3) >> 31) != 0;
        const uint32_t extent = Word(h, 4);
        if (extent < XTC_HEADER || pos + extent > end)
            return;
        const uint64_t payload = pos + XTC_HEADER, size = extent - XTC_HEADER;

        if (type == ID_XTC)
            Walk(f, file, payload, pos + extent, event, index);
        else if (type == ID_FRAME && event) {
            uint8_t fh[FRAME_HEADER];
            if (size >= FRAME_HEADER && ReadAt(f, payload, fh, FRAME_HEADER)) {
                const uint32_t width = Word(fh, 0), height = Word(fh, 1), depth = Word(fh, 2);
                const uint64_t npixel = static_cast<uint64_t>(width) * height;
                // Only a whole, undamaged 16-bit image of the version described above is read.
                if (depth > 8 && depth <= 16 && npixel > 0) {
                    if (damage != 0 || compressed || version != 1 || FRAME_HEADER + 2 * npixel > size)
                        index.damaged++;
                    else
                        index.frames.push_back({file, payload + FRAME_HEADER, width, height, Word(fh, 3), source});
                }
            }
        } else if (type == ID_EBEAM && event && version == 7 && !compressed && damage == 0 && size >= EBEAM_V7_SIZE) {
            uint8_t eb[EBEAM_V7_SIZE];
            double energy = 0;
            if (ReadAt(f, payload, eb, EBEAM_V7_SIZE) && Word(eb, 0) == 0) {
                std::memcpy(&energy, eb + EBEAM_V7_PHOTON_ENERGY, 8);
                if (energy > 100.0 && energy < 100000.0)
                    index.photon_energy_eV.push_back(energy);
            }
        } else if (type == ID_RAYONIX_CONFIG && version == 2 && !compressed && size >= 2 && size < 4096) {
            std::vector<uint8_t> cfg(size);
            if (ReadAt(f, payload, cfg.data(), size)) {
                index.rayonix_binning = cfg[0] == cfg[1] ? cfg[0] : 0;
                const std::string text(cfg.begin(), cfg.end());
                const auto mx = text.find("MX");
                if (mx != std::string::npos)
                    index.rayonix_model = text.substr(mx, text.find('\0', mx) - mx);
            }
        }
        pos += extent;
    }
}

}  // namespace

namespace xtc {

bool CanRead(const std::string &path) {
    std::error_code ec;
    if (!std::filesystem::is_regular_file(path, ec) || !path.ends_with(".xtc") || path.ends_with(".smd.xtc"))
        return false;
    std::ifstream f(path, std::ios::binary);
    uint8_t h[DGRAM_HEADER];
    return f && ReadAt(f, 0, h, DGRAM_HEADER) && PlausibleFirstDatagram(h);
}

Index IndexFile(const std::string &file) {
    std::ifstream f = Open(file);
    const uint64_t file_size = std::filesystem::file_size(file);
    Index index;
    uint64_t pos = 0;
    uint8_t h[DGRAM_HEADER];
    while (ReadAt(f, pos, h, DGRAM_HEADER)) {
        const uint32_t extent = Word(h, 9);
        if (extent < XTC_HEADER)
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                                  file + ": a datagram's extent is shorter than its own header");
        const uint64_t payload = pos + DGRAM_HEADER, end = pos + DGRAM_HEADER - XTC_HEADER + extent;
        // A partly downloaded or partly written stream ends inside its last datagram.
        if (end > file_size) {
            index.truncated = true;
            break;
        }
        Walk(f, file, payload, end, Transition(h) == TRANSITION_L1ACCEPT, index);
        pos = end;
    }
    return index;
}

void ReadFrame(const Frame &frame, int32_t *out, std::vector<uint8_t> &scratch) {
    const size_t npixel = static_cast<size_t>(frame.width) * frame.height;
    scratch.resize(2 * npixel);
    std::ifstream f = Open(frame.file);
    if (!ReadAt(f, frame.offset, scratch.data(), scratch.size()))
        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
                              frame.file + ": an image ends past the end of the file");
    for (size_t i = 0; i < npixel; i++) {
        uint16_t v;
        std::memcpy(&v, scratch.data() + 2 * i, 2);
        out[i] = static_cast<int32_t>(v) - static_cast<int32_t>(frame.pixel_offset);
    }
}

}  // namespace xtc