// SPDX-License-Identifier: MIT
// SPDX-FileCopyrightText: 2026 Martino Ferrari

#include "capture.hpp"

#include <fcntl.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <unistd.h>

#include <cstdio>
#include <cstring>
#include <ctime>

#include <chrono>
#include <vector>

namespace udpscope {
namespace {

/// The file header is one 64-byte block: fixed fields first, zeroed padding after.
constexpr std::uint64_t kHeaderBytes = 64;
/// `seq` lives at the *end* of a slot, at a fixed offset, because the reader must
/// find it without trusting any other field.
constexpr std::uint64_t kSeqBytes = 8;

void put_u32(std::uint8_t* p, std::uint32_t v) { std::memcpy(p, &v, 4); }
void put_u16(std::uint8_t* p, std::uint16_t v) { std::memcpy(p, &v, 2); }
void put_f64(std::uint8_t* p, double v) { std::memcpy(p, &v, 8); }
std::uint32_t get_u32(const std::uint8_t* p) {
  std::uint32_t v;
  std::memcpy(&v, p, 4);
  return v;
}
double get_f64(const std::uint8_t* p) {
  double v;
  std::memcpy(&v, p, 8);
  return v;
}

double wall_seconds() {
  struct timespec ts;
  clock_gettime(CLOCK_REALTIME, &ts);
  return double(ts.tv_sec) + double(ts.tv_nsec) * 1e-9;
}

/// `seq` is published with a release store *after* the bytes are down, so a reader
/// that sees a non-zero `seq` sees the whole slot -- the reason no `fsync` is needed
/// for the file to be readable after a crash. Plain __atomic builtins, not
/// `std::atomic_ref`: the target is mmap'd memory the class does not own as an object.
void publish_seq(std::uint8_t* slot, std::uint64_t seq) {
  __atomic_store_n(reinterpret_cast<std::uint64_t*>(slot + kSlotBytes - kSeqBytes), seq,
                   __ATOMIC_RELEASE);
}

std::uint64_t read_seq(const std::uint8_t* slot) {
  return __atomic_load_n(reinterpret_cast<const std::uint64_t*>(slot + kSlotBytes - kSeqBytes),
                         __ATOMIC_ACQUIRE);
}

}  // namespace

void PacketCapture::set_error(const char* what) noexcept {
  std::snprintf(error_, sizeof(error_), "%s", what);
  error_set_.store(true, std::memory_order_release);
}

void PacketCapture::error(char* out, std::size_t cap) const noexcept {
  if (cap == 0) return;
  out[0] = '\0';
  if (!error_set_.load(std::memory_order_acquire)) return;
  std::snprintf(out, cap, "%s", error_);
}

bool PacketCapture::open(const std::string& path, std::uint32_t slot_count) {
  close();
  path_ = path;
  slot_count_ = slot_count;
  seq_ = 0;
  cursor_ = 0;
  stored_.store(0, std::memory_order_relaxed);
  truncated_.store(0, std::memory_order_relaxed);
  error_set_.store(false, std::memory_order_relaxed);

  const std::chrono::steady_clock::time_point t0 = std::chrono::steady_clock::now();

  // Keep the previous capture one deep: re-enabling must not destroy the capture the
  // user enabled *for*. A failure here is reported, and the old previous kept.
  std::string prev = path;
  const std::size_t slash = prev.rfind('/');
  const std::size_t dot = prev.rfind('.');
  if (dot != std::string::npos && (slash == std::string::npos || dot > slash)) {
    prev.insert(dot, "-prev");
  } else {
    prev += "-prev";
  }
  ::rename(path.c_str(), prev.c_str());

  fd_ = ::open(path.c_str(), O_RDWR | O_CREAT | O_EXCL | O_CLOEXEC, 0600);
  if (fd_ < 0) {
    set_error(strerror(errno));
    return false;
  }

  // Reserved, not sparse: a store into a page the filesystem cannot back raises
  // SIGBUS and kills udpscope outright, and an out-of-space condition must be an
  // error at the moment the user clicked, not a crash ten minutes later.
  const std::uint64_t total = kHeaderBytes + std::uint64_t(slot_count) * kSlotBytes;
  if (::posix_fallocate(fd_, 0, off_t(total)) != 0) {
    set_error(strerror(errno));
    ::close(fd_);
    fd_ = -1;
    ::unlink(path.c_str());
    return false;
  }

  map_ = static_cast<std::uint8_t*>(
      ::mmap(nullptr, total, PROT_READ | PROT_WRITE, MAP_SHARED, fd_, 0));
  if (map_ == MAP_FAILED) {
    map_ = nullptr;
    set_error(strerror(errno));
    ::close(fd_);
    fd_ = -1;
    ::unlink(path.c_str());
    return false;
  }
  map_bytes_ = total;

  std::uint8_t* h = map_;
  std::memset(h, 0, kHeaderBytes);
  std::memcpy(h, "UDPSCAP\0", 8);
  put_u32(h + 8, 1);                      // version
  put_u32(h + 12, slot_count);
  put_u32(h + 16, kSlotBytes);
  put_u32(h + 20, 0);                     // cursor: a hint, `seq` is the order
  put_f64(h + 24, wall_seconds());        // the epoch recv_time is relative to
  put_u32(h + 32, peer_addr_);
  put_u16(h + 36, peer_port_);
  put_u16(h + 38, 0);  // local_port: the client does not expose its own socket
  const double ms =
      std::chrono::duration<double, std::milli>(std::chrono::steady_clock::now() - t0).count();
  open_ms_.store(ms, std::memory_order_relaxed);

  enabled_.store(true, std::memory_order_release);
  return true;
}

void PacketCapture::close() noexcept {
  if (map_ != nullptr) {
    ::munmap(map_, map_bytes_);
    map_ = nullptr;
  }
  if (fd_ >= 0) {
    ::close(fd_);
    fd_ = -1;
  }
  map_bytes_ = 0;
  enabled_.store(false, std::memory_order_relaxed);
}

void PacketCapture::store(const std::uint8_t* bytes, std::size_t len, double recv_time) noexcept {
  if (!enabled_.load(std::memory_order_relaxed) || map_ == nullptr) return;

  std::uint8_t* slot = map_ + kHeaderBytes + std::uint64_t(cursor_) * kSlotBytes;
  std::uint32_t orig = std::uint32_t(len);
  std::size_t fit = len;
  if (fit > kMaxDatagramBytes) {
    fit = kMaxDatagramBytes;
    truncated_.fetch_add(1, std::memory_order_relaxed);
  }

  put_f64(slot, recv_time);
  put_u32(slot + 8, std::uint32_t(fit));
  put_u32(slot + 12, orig);
  std::memcpy(slot + kRecordHeaderBytes, bytes, fit);
  ++seq_;
  publish_seq(slot, seq_);  // last: a non-zero `seq` is what makes a slot complete

  cursor_ = (cursor_ + 1) % slot_count_;
  put_u32(map_ + 20, cursor_);  // a hint for a human in xxd; `seq` is the order
  stored_.fetch_add(1, std::memory_order_relaxed);
}

bool read_capture(const std::string& path, CaptureReadStats* out) {
  if (out == nullptr) return false;
  *out = CaptureReadStats{};

  FILE* f = std::fopen(path.c_str(), "rb");
  if (f == nullptr) return false;
  std::fseek(f, 0, SEEK_END);
  const long size = std::ftell(f);
  std::fseek(f, 0, SEEK_SET);
  out->file_bytes = std::uint64_t(size < 0 ? 0 : size);

  std::uint8_t h[kHeaderBytes];
  if (std::fread(h, 1, sizeof(h), f) != sizeof(h)) {
    std::fclose(f);
    return false;
  }
  out->header_ok = std::memcmp(h, "UDPSCAP\0", 8) == 0 && get_u32(h + 8) == 1;
  out->slot_count = get_u32(h + 12);
  out->slot_bytes = get_u32(h + 16);
  out->cursor = get_u32(h + 20);
  out->t0_realtime = get_f64(h + 24);
  if (!out->header_ok || out->slot_count == 0 || out->slot_bytes < kRecordHeaderBytes + 8) {
    std::fclose(f);
    return false;
  }

  const std::uint32_t payload = out->slot_bytes - kRecordHeaderBytes - kSeqBytes;
  std::vector<std::uint8_t> slot(out->slot_bytes);
  for (std::uint32_t i = 0; i < out->slot_count; ++i) {
    if (std::fseek(f, long(kHeaderBytes + std::uint64_t(i) * out->slot_bytes), SEEK_SET) != 0) {
      break;
    }
    if (std::fread(slot.data(), 1, slot.size(), f) != slot.size()) break;
    const std::uint64_t seq = read_seq(slot.data());
    const std::uint32_t len = get_u32(slot.data() + 8);
    const std::uint32_t orig = get_u32(slot.data() + 12);
    if (seq == 0 && len == 0 && orig == 0) {
      ++out->empty;
      continue;
    }
    if (seq == 0) {
      ++out->torn;  // bytes down, ordinal never published: a capture that died mid-copy
      continue;
    }
    ++out->stored;
    out->payload += len;
    if (orig > payload) ++out->truncated;
    if (out->first_seq == 0 || seq < out->first_seq) out->first_seq = seq;
    if (seq > out->last_seq) out->last_seq = seq;
  }
  std::fclose(f);
  return true;
}

}  // namespace udpscope