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

#include "source/e1608/e1608_link.hpp"

#include <algorithm>
#include <arpa/inet.h>
#include <cerrno>
#include <chrono>
#include <cstdio>
#include <cstring>
#include <ifaddrs.h>
#include <net/if.h>
#include <netdb.h>
#include <netinet/in.h>
#include <netinet/tcp.h>
#include <poll.h>
#include <sys/socket.h>
#include <thread>
#include <unistd.h>

namespace udpscope::e1608 {

void Fd::reset(int fd) noexcept {
  if (fd_ >= 0) ::close(fd_);
  fd_ = fd;
}

namespace {

using Clock = std::chrono::steady_clock;
using Kind = LinkError::Kind;

constexpr std::size_t kMaxFrame = 1024;

[[noreturn]] void sys_fail(const std::string& what) {
  throw LinkError(Kind::Io, what + ": " + std::strerror(errno));
}

int ms_left(Clock::time_point deadline) {
  const auto d = std::chrono::duration_cast<std::chrono::milliseconds>(deadline - Clock::now());
  return d.count() < 0 ? 0 : static_cast<int>(d.count());
}

/// True when `fd` is ready for `events` (or has an error/hang-up to report); false on timeout.
bool wait_fd(int fd, short events, int timeout_ms) {
  pollfd p{fd, events, 0};
  for (;;) {
    const int r = ::poll(&p, 1, timeout_ms);
    if (r > 0) return true;
    if (r == 0) return false;
    if (errno == EINTR) continue;
    sys_fail("poll");
  }
}

sockaddr_in make_addr(const std::string& host, int port) {
  sockaddr_in a{};
  a.sin_family = AF_INET;
  a.sin_port = htons(static_cast<std::uint16_t>(port));
  if (::inet_pton(AF_INET, host.c_str(), &a.sin_addr) == 1) return a;
  addrinfo hints{};
  hints.ai_family = AF_INET;
  hints.ai_socktype = SOCK_STREAM;
  addrinfo* res = nullptr;
  if (::getaddrinfo(host.c_str(), nullptr, &hints, &res) != 0 || res == nullptr)
    throw LinkError(Kind::Io, "cannot resolve '" + host + "'");
  a.sin_addr = reinterpret_cast<const sockaddr_in*>(res->ai_addr)->sin_addr;
  ::freeaddrinfo(res);
  return a;
}

void send_all(int fd, const std::uint8_t* p, std::size_t n, Clock::time_point deadline) {
  while (n > 0) {
    if (!wait_fd(fd, POLLOUT, ms_left(deadline)))
      throw LinkError(Kind::Timeout, "timed out sending to the board");
    const ssize_t w = ::send(fd, p, n, MSG_NOSIGNAL | MSG_DONTWAIT);
    if (w < 0) {
      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) continue;
      sys_fail("send");
    }
    p += w;
    n -= static_cast<std::size_t>(w);
  }
}

void read_exact(int fd, std::uint8_t* p, std::size_t n, Clock::time_point deadline) {
  while (n > 0) {
    if (!wait_fd(fd, POLLIN, ms_left(deadline)))
      throw LinkError(Kind::Timeout, "timed out waiting for the board");
    const ssize_t r = ::recv(fd, p, n, MSG_DONTWAIT);
    if (r == 0) throw LinkError(Kind::Io, "the board closed the connection");
    if (r < 0) {
      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) continue;
      sys_fail("recv");
    }
    p += r;
    n -= static_cast<std::size_t>(r);
  }
}

/// A connected, non-blocking TCP socket with TCP_NODELAY. `rcvbuf > 0` sets SO_RCVBUF first.
Fd tcp_connect(const std::string& host, int port, int timeout_ms, int rcvbuf) {
  const sockaddr_in a = make_addr(host, port);
  Fd fd(::socket(AF_INET, SOCK_STREAM | SOCK_NONBLOCK, 0));
  if (!fd.valid()) sys_fail("socket");
  const int one = 1;
  ::setsockopt(fd.get(), IPPROTO_TCP, TCP_NODELAY, &one, sizeof one);
  if (rcvbuf > 0) ::setsockopt(fd.get(), SOL_SOCKET, SO_RCVBUF, &rcvbuf, sizeof rcvbuf);
  const std::string where = host + ":" + std::to_string(port);
  if (::connect(fd.get(), reinterpret_cast<const sockaddr*>(&a), sizeof a) < 0 &&
      errno != EINPROGRESS)
    sys_fail("cannot connect to " + where);
  if (!wait_fd(fd.get(), POLLOUT, timeout_ms))
    throw LinkError(Kind::Timeout, "timed out connecting to " + where);
  int err = 0;
  socklen_t len = sizeof err;
  ::getsockopt(fd.get(), SOL_SOCKET, SO_ERROR, &err, &len);
  if (err != 0) {
    errno = err;
    sys_fail("cannot connect to " + where);
  }
  return fd;
}

/// A `D` reply (at least 29 bytes, first byte `D`, our product id). True when it added a board
/// that `found` did not already hold (one entry per MAC).
bool take_reply(const std::uint8_t* buf, ssize_t n, const sockaddr_in& from,
                std::vector<DeviceInfo>& found) {
  if (n < 29 || buf[0] != 'D') return false;
  DeviceInfo info;
  info.product_id = static_cast<std::uint16_t>(buf[7] | (buf[8] << 8));
  if (info.product_id != kProductId) return false;
  std::copy(buf + 1, buf + 7, info.mac.begin());
  for (const DeviceInfo& d : found) {
    if (d.mac == info.mac) return false;
  }
  char ip[INET_ADDRSTRLEN] = {};
  ::inet_ntop(AF_INET, &from.sin_addr, ip, sizeof ip);
  info.ip = ip;
  found.push_back(info);
  return true;
}

}  // namespace

// --- addresses ------------------------------------------------------------------------------

bool parse_mac(const std::string& text, Mac& out) {
  Mac m{};
  std::size_t pos = 0;
  for (std::size_t i = 0; i < 6; ++i) {
    if (pos + 2 > text.size()) return false;
    unsigned v = 0;
    for (std::size_t k = 0; k < 2; ++k) {
      const char c = text[pos + k];
      v <<= 4;
      if (c >= '0' && c <= '9') v |= static_cast<unsigned>(c - '0');
      else if (c >= 'a' && c <= 'f') v |= static_cast<unsigned>(c - 'a' + 10);
      else if (c >= 'A' && c <= 'F') v |= static_cast<unsigned>(c - 'A' + 10);
      else return false;
    }
    m[i] = static_cast<std::uint8_t>(v);
    pos += 2;
    if (i < 5) {
      if (pos >= text.size() || (text[pos] != ':' && text[pos] != '-')) return false;
      ++pos;
    }
  }
  if (pos != text.size()) return false;
  out = m;
  return true;
}

std::string format_mac(const Mac& mac) {
  char buf[24];
  std::snprintf(buf, sizeof buf, "%02x:%02x:%02x:%02x:%02x:%02x", mac[0], mac[1], mac[2], mac[3],
                mac[4], mac[5]);
  return buf;
}

std::vector<std::string> sweep_targets(const std::string& subnet) {
  std::string s = subnet;
  while (!s.empty() && s.front() == ' ') s.erase(s.begin());
  while (!s.empty() && (s.back() == ' ' || s.back() == '.')) s.pop_back();
  if (s.empty()) return {"255.255.255.255"};
  int parts = 0;
  std::size_t i = 0;
  while (i <= s.size()) {
    const std::size_t dot = s.find('.', i);
    const std::string part = s.substr(i, dot == std::string::npos ? std::string::npos : dot - i);
    if (part.empty() || part.size() > 3) return {};
    int v = 0;
    for (char c : part) {
      if (c < '0' || c > '9') return {};
      v = v * 10 + (c - '0');
    }
    if (v > 255) return {};
    ++parts;
    if (dot == std::string::npos) break;
    i = dot + 1;
  }
  if (parts != 3) return {};
  std::vector<std::string> out;
  out.reserve(254);
  for (int h = 1; h <= 254; ++h) out.push_back(s + "." + std::to_string(h));
  return out;
}

std::vector<DeviceInfo> discover(const std::vector<std::string>& targets, int port, int wait_ms) {
  std::vector<DeviceInfo> found;
  Fd u(::socket(AF_INET, SOCK_DGRAM | SOCK_NONBLOCK, 0));
  if (!u.valid()) sys_fail("socket");
  const int one = 1;
  ::setsockopt(u.get(), SOL_SOCKET, SO_BROADCAST, &one, sizeof one);
  const std::uint8_t probe = 'D';
  for (const std::string& t : targets) {
    sockaddr_in a{};
    try {
      a = make_addr(t, port);
    } catch (const LinkError&) {
      continue;
    }
    // A sweep may hit unreachable hosts; those sends fail and are simply skipped. A full send
    // buffer is not that: wait for room once and retry, so a large sweep does not lose probes.
    const auto send_probe = [&] {
      return ::sendto(u.get(), &probe, 1, MSG_NOSIGNAL, reinterpret_cast<const sockaddr*>(&a),
                      sizeof a);
    };
    if (send_probe() < 0 && (errno == EAGAIN || errno == ENOBUFS)) {
      wait_fd(u.get(), POLLOUT, 20);
      (void)send_probe();
    }
  }
  const auto deadline = Clock::now() + std::chrono::milliseconds(wait_ms);
  for (;;) {
    if (!wait_fd(u.get(), POLLIN, ms_left(deadline))) break;
    std::uint8_t buf[128];
    sockaddr_in from{};
    socklen_t flen = sizeof from;
    const ssize_t n = ::recvfrom(u.get(), buf, sizeof buf, MSG_DONTWAIT,
                                 reinterpret_cast<sockaddr*>(&from), &flen);
    if (n < 0) {
      if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) continue;
      break;
    }
    take_reply(buf, n, from, found);
  }
  return found;
}

Scan::Scan(int port) {
  ifaddrs* list = nullptr;
  if (::getifaddrs(&list) != 0) return;
  sockaddr_in to{};
  to.sin_family = AF_INET;
  to.sin_port = htons(static_cast<std::uint16_t>(port));
  to.sin_addr.s_addr = htonl(INADDR_BROADCAST);
  const std::uint8_t probe = 'D';
  std::vector<std::string> done;  // an interface with several addresses needs one probe, not several
  for (const ifaddrs* a = list; a != nullptr; a = a->ifa_next) {
    if (a->ifa_addr == nullptr || a->ifa_addr->sa_family != AF_INET) continue;
    const unsigned fl = a->ifa_flags;
    if ((fl & IFF_UP) == 0 || (fl & IFF_LOOPBACK) != 0 || (fl & IFF_BROADCAST) == 0) continue;
    const std::string name = a->ifa_name;
    if (std::find(done.begin(), done.end(), name) != done.end()) continue;
    Fd s(::socket(AF_INET, SOCK_DGRAM | SOCK_NONBLOCK, 0));
    if (!s.valid()) continue;
    const int one = 1;
    ::setsockopt(s.get(), SOL_SOCKET, SO_BROADCAST, &one, sizeof one);
    if (::setsockopt(s.get(), SOL_SOCKET, SO_BINDTODEVICE, name.c_str(),
                     static_cast<socklen_t>(name.size())) != 0) {
      sockaddr_in me = *reinterpret_cast<const sockaddr_in*>(a->ifa_addr);
      me.sin_port = 0;
      if (::bind(s.get(), reinterpret_cast<const sockaddr*>(&me), sizeof me) != 0) continue;
    }
    if (::sendto(s.get(), &probe, 1, MSG_NOSIGNAL, reinterpret_cast<const sockaddr*>(&to),
                 sizeof to) != 1)
      continue;
    done.push_back(name);
    socks_.push_back(std::move(s));
  }
  ::freeifaddrs(list);
}

const std::vector<DeviceInfo>& Scan::poll() {
  for (const Fd& s : socks_) {
    for (;;) {
      std::uint8_t buf[128];
      sockaddr_in from{};
      socklen_t flen = sizeof from;
      const ssize_t n = ::recvfrom(s.get(), buf, sizeof buf, MSG_DONTWAIT,
                                   reinterpret_cast<sockaddr*>(&from), &flen);
      if (n < 0) {
        if (errno == EINTR) continue;
        break;  // nothing more on this socket for now
      }
      take_reply(buf, n, from, found_);
    }
  }
  return found_;
}

std::vector<DeviceInfo> discover_broadcast(int port, int wait_ms) {
  Scan scan(port);
  const auto deadline = Clock::now() + std::chrono::milliseconds(wait_ms);
  while (ms_left(deadline) > 0) {
    scan.poll();
    std::this_thread::sleep_for(std::chrono::milliseconds(10));
  }
  return scan.poll();
}

// --- Link -----------------------------------------------------------------------------------

void Link::connect(const std::string& host, int port, std::uint32_t code, int timeout_ms) {
  close();
  host_ = host;
  port_ = port;
  timeout_ms_ = timeout_ms;
  next_id_ = 0;

  // The connection-code handshake: a connected UDP socket, so an ICMP refusal surfaces.
  {
    const sockaddr_in a = make_addr(host, port);
    Fd u(::socket(AF_INET, SOCK_DGRAM, 0));
    if (!u.valid()) sys_fail("socket");
    if (::connect(u.get(), reinterpret_cast<const sockaddr*>(&a), sizeof a) < 0)
      sys_fail("handshake with " + host);
    const std::uint8_t msg[5] = {'C', static_cast<std::uint8_t>(code & 0xFF),
                                 static_cast<std::uint8_t>((code >> 8) & 0xFF),
                                 static_cast<std::uint8_t>((code >> 16) & 0xFF),
                                 static_cast<std::uint8_t>((code >> 24) & 0xFF)};
    if (::send(u.get(), msg, sizeof msg, MSG_NOSIGNAL) < 0) sys_fail("handshake with " + host);
    if (!wait_fd(u.get(), POLLIN, timeout_ms))
      throw LinkError(Kind::Timeout, "no answer from the board at " + host);
    std::uint8_t rep[16];
    const ssize_t n = ::recv(u.get(), rep, sizeof rep, 0);
    if (n < 0) sys_fail("handshake with " + host);
    if (n < 2 || rep[0] != 'C') throw LinkError(Kind::Protocol, "unexpected handshake reply");
    switch (rep[1]) {
      case 0:
      case 2:
        break;
      case 1:
        throw LinkError(Kind::Rejected, "the board rejected the connection code");
      case 3:
        throw LinkError(Kind::Busy, "the board is busy (another client holds it)");
      default:
        throw LinkError(Kind::Protocol, "unexpected handshake status " + std::to_string(rep[1]));
    }
  }
  cmd_ = tcp_connect(host, port, timeout_ms, 0);
}

void Link::close() {
  scan_.reset();
  cmd_.reset();
}

std::vector<std::uint8_t> Link::command(std::uint8_t cmd, const std::vector<std::uint8_t>& payload,
                                        int expect) {
  if (!cmd_.valid()) throw LinkError(Kind::Io, "not connected");
  if (payload.size() + 7 > kMaxFrame) throw LinkError(Kind::Protocol, "command too long");
  // Bytes left by an earlier exchange that timed out must not be read as this reply.
  {
    std::uint8_t junk[256];
    while (::recv(cmd_.get(), junk, sizeof junk, MSG_DONTWAIT) > 0) {
    }
  }
  if (++next_id_ == 0) next_id_ = 1;  // id 0 is never used, in case a board treats it specially
  const std::uint8_t id = next_id_;
  std::vector<std::uint8_t> f;
  f.reserve(7 + payload.size());
  f.push_back(0xDB);
  f.push_back(cmd);
  f.push_back(id);
  f.push_back(0);
  f.push_back(static_cast<std::uint8_t>(payload.size() & 0xFF));
  f.push_back(static_cast<std::uint8_t>((payload.size() >> 8) & 0xFF));
  f.insert(f.end(), payload.begin(), payload.end());
  unsigned sum = 0;
  for (std::uint8_t b : f) sum += b;
  f.push_back(static_cast<std::uint8_t>(0xFF - (sum & 0xFF)));

  const auto deadline = Clock::now() + std::chrono::milliseconds(timeout_ms_);
  send_all(cmd_.get(), f.data(), f.size(), deadline);

  std::uint8_t h[6];
  read_exact(cmd_.get(), h, sizeof h, deadline);
  if (h[0] != 0xDB) throw LinkError(Kind::Protocol, "bad sync byte in a reply");
  const std::size_t n = static_cast<std::size_t>(h[4] | (h[5] << 8));
  if (n + 7 > kMaxFrame) throw LinkError(Kind::Protocol, "oversized reply");
  std::vector<std::uint8_t> body(n + 1);
  read_exact(cmd_.get(), body.data(), body.size(), deadline);
  unsigned rsum = 0;
  for (std::uint8_t b : h) rsum += b;
  for (std::uint8_t b : body) rsum += b;
  if ((rsum & 0xFF) != 0xFF) throw LinkError(Kind::Protocol, "bad checksum in a reply");
  if (h[1] != (cmd | 0x80) || h[2] != id)
    throw LinkError(Kind::Protocol, "reply does not match the command");
  if (h[3] != 0) {
    static const char* const kText[] = {"ok",        "protocol error", "parameter error",
                                        "busy",      "not ready",      "timeout",
                                        "other error"};
    const char* t = h[3] < 7 ? kText[h[3]] : "unknown error";
    char hex[8];
    std::snprintf(hex, sizeof hex, "%02x", cmd);
    throw LinkError(Kind::Command, std::string("the board refused command 0x") + hex + ": " + t);
  }
  body.pop_back();
  if (expect >= 0 && body.size() != static_cast<std::size_t>(expect))
    throw LinkError(Kind::Command, "reply of " + std::to_string(body.size()) +
                                       " bytes, expected " + std::to_string(expect));
  return body;
}

void Link::open_scan() {
  close_scan();
  scan_ = tcp_connect(host_, port_ + 1, timeout_ms_, 1 << 20);
}

void Link::close_scan() { scan_.reset(); }

void Link::send_filler() {
  if (!cmd_.valid()) throw LinkError(Kind::Io, "not connected");
  const std::uint8_t z = 0;
  send_all(cmd_.get(), &z, 1, Clock::now() + std::chrono::milliseconds(timeout_ms_));
}

std::size_t Link::read_scan(std::uint8_t* buf, std::size_t cap, int timeout_ms) {
  if (!scan_.valid()) throw LinkError(Kind::Io, "the scan socket is not open");
  if (cap == 0) return 0;  // recv would return 0 and read as a close
  if (!wait_fd(scan_.get(), POLLIN, timeout_ms)) return 0;
  const ssize_t n = ::recv(scan_.get(), buf, cap, MSG_DONTWAIT);
  if (n > 0) return static_cast<std::size_t>(n);
  if (n == 0) throw LinkError(Kind::Io, "scan socket closed by the board");
  if (errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) return 0;
  sys_fail("scan recv");
}

}  // namespace udpscope::e1608