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

#include "source/flatudp/flatudp_source.hpp"

#include <arpa/inet.h>
#include <netinet/in.h>
#include <poll.h>
#include <sys/socket.h>
#include <unistd.h>

#include <algorithm>
#include <cctype>
#include <cerrno>
#include <cmath>
#include <cstdio>
#include <cstring>
#include <stdexcept>
#include <thread>

namespace udpscope::flatudp {
namespace {

using namespace std::chrono_literals;

std::string lower(std::string s) {
  for (char& c : s) c = static_cast<char>(std::tolower(static_cast<unsigned char>(c)));
  return s;
}

std::string trim(const std::string& s) {
  std::size_t a = 0, b = s.size();
  while (a < b && std::isspace(static_cast<unsigned char>(s[a]))) ++a;
  while (b > a && std::isspace(static_cast<unsigned char>(s[b - 1]))) --b;
  return s.substr(a, b - a);
}

double epoch_seconds() {
  return std::chrono::duration<double>(std::chrono::system_clock::now().time_since_epoch()).count();
}

constexpr std::size_t kMaxReportedSizes = 8;  ///< a sender of random sizes must not flood stderr

}  // namespace

bool valid_ipv4(const std::string& s) {
  in_addr a{};
  return inet_pton(AF_INET, s.c_str(), &a) == 1;
}

bool valid_multicast(const std::string& s) {
  in_addr a{};
  if (inet_pton(AF_INET, s.c_str(), &a) != 1) return false;
  return (ntohl(a.s_addr) >> 28) == 0xE;
}

// --- FlatUdpSettings ------------------------------------------------------------------------

FlatUdpSettings FlatUdpSettings::from(const SourceConfig& c) {
  FlatUdpSettings s;
  s.port = static_cast<int>(c.get_int("port", kDefaultPort));
  s.bind = trim(c.get("bind"));
  s.multicast = trim(c.get("multicast"));
  s.iface = trim(c.get("iface"));
  s.rate = c.get_double("rate", 0.0);
  const std::string e = lower(trim(c.get("endian", "little")));
  s.endian_ok = e == "little" || e == "big";
  s.big_endian = e == "big";
  s.signals_error = parse(c.get("signals"), s.rows);
  return s;
}

std::string FlatUdpSettings::problem() const {
  if (port < 1 || port > 65535) return "the port must be 1 to 65535";
  if (!signals_error.empty()) return signals_error;
  const std::string p = flatudp::problem(rows);
  if (!p.empty()) return p;
  if (!(rate > 0.0) || !std::isfinite(rate)) return "give the sender's cycle rate in Hz (rate)";
  if (!endian_ok) return "endian must be little or big";
  if (!bind.empty() && !valid_ipv4(bind)) return "the bind address '" + bind + "' is not an IPv4 address";
  if (!multicast.empty() && !valid_multicast(multicast))
    return "'" + multicast + "' is not a multicast group (224.0.0.0 to 239.255.255.255)";
  if (!iface.empty() && !valid_ipv4(iface))
    return "the interface address '" + iface + "' is not an IPv4 address";
  return {};
}

std::string FlatUdpSettings::identity() const {
  return (multicast.empty() ? std::string() : multicast + ":") + std::to_string(port);
}

std::size_t FlatUdpSettings::block() const {
  if (!all_scalar(rows) || !(rate > 0.0) || !std::isfinite(rate)) return 1;
  const long b = std::lround(rate / 100.0);
  return static_cast<std::size_t>(std::min(1024L, std::max(1L, b)));
}

// --- FlatUdpSource --------------------------------------------------------------------------

FlatUdpSource::FlatUdpSource(const SourceConfig& cfg) : cfg_(FlatUdpSettings::from(cfg)) {}

FlatUdpSource::~FlatUdpSource() { close_fd(); }

std::string FlatUdpSource::last_message() const {
  std::lock_guard<std::mutex> lk(msg_mu_);
  return last_message_;
}

void FlatUdpSource::note(const std::string& msg) {
  {
    std::lock_guard<std::mutex> lk(msg_mu_);
    last_message_ = msg;
  }
  std::fprintf(stderr, "[flatudp] %s\n", msg.c_str());
  std::fflush(stderr);
}

void FlatUdpSource::close_fd() {
  if (fd_ >= 0) {
    ::close(fd_);
    fd_ = -1;
  }
}

void FlatUdpSource::open(SourceSink& sink) {
  const std::string p = cfg_.problem();
  if (!p.empty()) throw std::runtime_error(p);
  sink_ = &sink;

  signals_.clear();
  for (const Row& r : cfg_.rows) {
    SignalInfo s;
    s.name = r.name;
    s.type_code = static_cast<std::uint8_t>(r.type);
    s.rows = r.count;
    s.cols = 1;
    s.time_mode = kTimePacket;
    s.sampling_rate = cfg_.rate;
    signals_.push_back(s);
  }
  expected_ = datagram_size(cfg_.rows);
  block_ = cfg_.block();
  pending_.assign(cfg_.rows.size(), {});
  pending_n_ = 0;
  rx_.assign(65536, 0);

  const auto fail = [this](const std::string& what) {
    const std::string msg = what + ": " + std::strerror(errno);
    close_fd();
    throw std::runtime_error(msg);
  };

  fd_ = ::socket(AF_INET, SOCK_DGRAM | SOCK_NONBLOCK | SOCK_CLOEXEC, 0);
  if (fd_ < 0) fail("cannot create a UDP socket");
  // Receivers of one multicast group share its port; a unicast port is the one app's.
  if (!cfg_.multicast.empty()) {
    const int one = 1;
    ::setsockopt(fd_, SOL_SOCKET, SO_REUSEADDR, &one, sizeof one);
  }
  const int rcvbuf = 4 << 20;  // best effort: a 1 kHz sender at 64 KiB would not fit a default one
  ::setsockopt(fd_, SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof rcvbuf);

  sockaddr_in addr{};
  addr.sin_family = AF_INET;
  addr.sin_port = htons(static_cast<std::uint16_t>(cfg_.port));
  addr.sin_addr.s_addr = htonl(INADDR_ANY);
  if (!cfg_.bind.empty()) inet_pton(AF_INET, cfg_.bind.c_str(), &addr.sin_addr);
  if (::bind(fd_, reinterpret_cast<const sockaddr*>(&addr), sizeof addr) < 0)
    fail("cannot bind UDP port " + std::to_string(cfg_.port));

  if (!cfg_.multicast.empty()) {
    ip_mreq mreq{};
    inet_pton(AF_INET, cfg_.multicast.c_str(), &mreq.imr_multiaddr);
    if (cfg_.iface.empty()) {
      mreq.imr_interface.s_addr = htonl(INADDR_ANY);
    } else {
      inet_pton(AF_INET, cfg_.iface.c_str(), &mreq.imr_interface);
    }
    if (::setsockopt(fd_, IPPROTO_IP, IP_ADD_MEMBERSHIP, &mreq, sizeof mreq) < 0)
      fail("cannot join multicast group " + cfg_.multicast);
  }
}

void FlatUdpSource::close() {
  close_fd();
  connected_.store(false, std::memory_order_relaxed);
}

SourceStats FlatUdpSource::stats() const {
  SourceStats s;
  s.packets = valid_.load(std::memory_order_relaxed) + wrong_.load(std::memory_order_relaxed);
  s.frames = frames_;
  s.stale = wrong_.load(std::memory_order_relaxed);
  s.reconnects = reconnects_;
  s.connected = connected_.load(std::memory_order_relaxed);
  return s;
}

void FlatUdpSource::poll(int timeout_ms) {
  drain_commands();
  // Never wait long: the partial-block flush and the silence timeout are checked on this cadence.
  const int wait = std::clamp(timeout_ms, 0, 20);
  if (fd_ < 0 || sink_ == nullptr) {
    std::this_thread::sleep_for(std::chrono::milliseconds(wait));
    return;
  }
  if (!configured_) {
    configured_ = true;
    sink_->on_config(signals_, 1);
  }
  pollfd p{};
  p.fd = fd_;
  p.events = POLLIN;
  if (::poll(&p, 1, wait) > 0 && (p.revents & POLLIN)) drain_socket();

  const auto now = Clock::now();
  if (pending_n_ > 0 && now - batch_start_ >= 100ms) flush();  // a stopped sender holds nothing back
  if (connected_.load(std::memory_order_relaxed) && now - last_valid_ >= 2s) {
    connected_.store(false, std::memory_order_relaxed);
    sink_->on_event(SourceEvent::Disconnected, "no datagram for 2 s");
  }
}

void FlatUdpSource::drain_socket() {
  for (int i = 0; i < 8192; ++i) {  // bounded, so a flood cannot starve the timers
    const ssize_t n = ::recv(fd_, rx_.data(), rx_.size(), 0);
    if (n < 0) {
      if (errno == EINTR) continue;
      break;  // EAGAIN: drained
    }
    const double t = epoch_seconds();
    const std::size_t len = static_cast<std::size_t>(n);
    if (tap_) tap_(rx_.data(), len, t);
    if (len == expected_) {
      accept(rx_.data());
    } else {
      reject(len);
    }
  }
}

void FlatUdpSource::accept(const std::uint8_t* data) {
  valid_.fetch_add(1, std::memory_order_relaxed);
  last_valid_ = Clock::now();
  if (!connected_.load(std::memory_order_relaxed)) {
    if (ever_connected_) ++reconnects_;
    ever_connected_ = true;
    connected_.store(true, std::memory_order_relaxed);
    sink_->on_event(SourceEvent::Connected, "");
  }
  decode_datagram(cfg_.rows, data, cfg_.big_endian, pending_);
  if (pending_n_ == 0) batch_start_ = last_valid_;
  if (++pending_n_ >= block_) flush();
}

void FlatUdpSource::reject(std::size_t len) {
  wrong_.fetch_add(1, std::memory_order_relaxed);
  if (reported_sizes_.size() >= kMaxReportedSizes || !reported_sizes_.insert(len).second) return;
  const std::string msg = "datagram of " + std::to_string(len) + " bytes, expected " +
                          std::to_string(expected_) + ", check the signal list";
  note(msg);
  sink_->on_event(SourceEvent::Error, msg.c_str());
}

void FlatUdpSource::flush() {
  const std::size_t n = signals_.size();
  views_.resize(n);
  for (std::size_t i = 0; i < n; ++i)
    views_[i] = SignalValues{pending_[i].data(), static_cast<std::uint32_t>(pending_[i].size())};
  Frame f;
  f.counter = frame_counter_++;
  f.lost = 0;
  f.late = false;
  f.recv_time = epoch_seconds();
  f.publish_mode = 1;
  f.num_samples = static_cast<std::uint32_t>(pending_n_);
  f.num_signals = static_cast<std::uint32_t>(n);
  f.signals = signals_.data();
  f.values = views_.data();
  ++frames_;
  sink_->on_data(f);
  for (auto& c : pending_) c.clear();
  pending_n_ = 0;
}

}  // namespace udpscope::flatudp