/*
 * This file is part of the ZoneMinder Project. See AUTHORS file for Copyright information
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License as published by the
 * Free Software Foundation; either version 2 of the License, or (at your
 * option) any later version.
 *
 * This program is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License for
 * more details.
 *
 * You should have received a copy of the GNU General Public License along
 * with this program. If not, see <http://www.gnu.org/licenses/>.
 */

#include "zm_catch2.h"

#include "zm_stream_socket.h"
#include "zm_stream_socket_client.h"
#include "zm_stream_socket_protocol.h"

#include <atomic>
#include <chrono>
#include <condition_variable>
#include <cstring>
#include <memory>
#include <mutex>
#include <string>
#include <sys/select.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <thread>
#include <unistd.h>
#include <vector>

using namespace zm::stream_socket;

namespace {

// Per-process path so concurrent test runs do not share a socket file
const std::string kSockPathStr =
    "/tmp/zm.stream_socket_client.unittest." + std::to_string(getpid()) + ".sock";
const char *kSockPath = kSockPathStr.c_str();

struct AVCodecParametersDeleter {
  void operator()(AVCodecParameters *par) const { avcodec_parameters_free(&par); }
};
using codec_parameters_ptr = std::unique_ptr<AVCodecParameters, AVCodecParametersDeleter>;

// Collects callback invocations with waiting helpers
struct Collector {
  std::mutex mutex;
  std::condition_variable cv;
  std::vector<std::pair<StreamId, HelloInfo>> hellos;
  std::vector<std::pair<Header, std::vector<uint8_t>>> media;
  std::vector<std::pair<Header, MonitorEvent>> events;
  int byes = 0;
  int disconnects = 0;

  StreamSocketClient::Callbacks MakeCallbacks() {
    StreamSocketClient::Callbacks callbacks;
    callbacks.on_hello = [this](StreamId stream, const HelloInfo &info, uint32_t) {
      std::lock_guard<std::mutex> lock(mutex);
      hellos.emplace_back(stream, info);
      cv.notify_all();
    };
    callbacks.on_event = [this](const Header &header, const MonitorEvent &event) {
      std::lock_guard<std::mutex> lock(mutex);
      events.emplace_back(header, event);
      cv.notify_all();
    };
    callbacks.on_media = [this](const Header &header, const uint8_t *data, size_t size) {
      std::lock_guard<std::mutex> lock(mutex);
      media.emplace_back(header, std::vector<uint8_t>(data, data + size));
      cv.notify_all();
    };
    callbacks.on_bye = [this]() {
      std::lock_guard<std::mutex> lock(mutex);
      ++byes;
      cv.notify_all();
    };
    callbacks.on_disconnect = [this]() {
      std::lock_guard<std::mutex> lock(mutex);
      ++disconnects;
      cv.notify_all();
    };
    return callbacks;
  }

  template <typename Pred>
  bool WaitFor(Pred pred, std::chrono::milliseconds timeout = std::chrono::seconds(5)) {
    std::unique_lock<std::mutex> lock(mutex);
    return cv.wait_for(lock, timeout, pred);
  }
};

av_packet_ptr make_packet(size_t size, uint8_t fill) {
  av_packet_ptr packet{av_packet_alloc()};
  REQUIRE(av_new_packet(packet.get(), size) == 0);
  memset(packet->data, fill, size);
  return packet;
}

}  // namespace

TEST_CASE("StreamSocketClient receives HELLO and media end to end", "[stream_socket_client]") {
  StreamSocket server(7, kSockPath);
  REQUIRE(server.Start());

  codec_parameters_ptr par{avcodec_parameters_alloc()};
  par->codec_type = AVMEDIA_TYPE_VIDEO;
  par->codec_id = AV_CODEC_ID_HEVC;
  par->width = 2560;
  par->height = 1440;
  server.SetVideoParams(par.get(), {20, 1});

  Collector collector;
  StreamSocketClient client(std::string(kSockPath), collector.MakeCallbacks());

  REQUIRE(collector.WaitFor([&] { return !collector.hellos.empty(); }));
  {
    std::lock_guard<std::mutex> lock(collector.mutex);
    REQUIRE(collector.hellos[0].first == StreamId::Video);
    REQUIRE(collector.hellos[0].second.codec_id == AV_CODEC_ID_HEVC);
    REQUIRE(collector.hellos[0].second.width == 2560);
    REQUIRE(collector.hellos[0].second.fps_num == 20);
  }

  // wait until the server has the client registered, then send media
  std::this_thread::sleep_for(std::chrono::milliseconds(100));
  av_packet_ptr packet = make_packet(4096, 0xC3);
  server.SendMedia(packet.get(), StreamId::Video, true, 777000);

  REQUIRE(collector.WaitFor([&] { return !collector.media.empty(); }));
  {
    std::lock_guard<std::mutex> lock(collector.mutex);
    const Header &header = collector.media[0].first;
    REQUIRE(header.type == static_cast<uint8_t>(MessageType::Media));
    REQUIRE(header.flags == kFlagKeyframe);
    REQUIRE(header.pts_us == 777000);
    REQUIRE(collector.media[0].second.size() == 4096);
    REQUIRE(collector.media[0].second[100] == 0xC3);
  }

  client.Stop();
  server.Stop();
}

TEST_CASE("StreamSocketClient gets BYE on server stop and reconnects", "[stream_socket_client]") {
  auto server = std::make_unique<StreamSocket>(7, kSockPath);
  REQUIRE(server->Start());

  codec_parameters_ptr par{avcodec_parameters_alloc()};
  par->codec_type = AVMEDIA_TYPE_VIDEO;
  par->codec_id = AV_CODEC_ID_H264;
  server->SetVideoParams(par.get(), {0, 0});

  Collector collector;
  StreamSocketClient client(std::string(kSockPath), collector.MakeCallbacks());

  REQUIRE(collector.WaitFor([&] { return collector.hellos.size() >= 1; }));

  server->Stop();
  REQUIRE(collector.WaitFor([&] { return collector.byes >= 1 && collector.disconnects >= 1; }));

  // Producer comes back (zmc restart): client reconnects and gets a new HELLO
  server = std::make_unique<StreamSocket>(7, kSockPath);
  REQUIRE(server->Start());
  server->SetVideoParams(par.get(), {0, 0});

  REQUIRE(collector.WaitFor([&] { return collector.hellos.size() >= 2; },
                            std::chrono::seconds(10)));

  client.Stop();
  server->Stop();
}

TEST_CASE("StreamSocketClient handles fragmented delivery", "[stream_socket_client]") {
  int fds[2];
  REQUIRE(socketpair(AF_UNIX, SOCK_STREAM, 0, fds) == 0);

  Collector collector;
  StreamSocketClient client(fds[0], collector.MakeCallbacks());

  // Build one MEDIA message and dribble it through in small fragments
  std::vector<uint8_t> payload(300, 0x5e);
  Header header = {};
  header.length = kHeaderLengthBytes + payload.size();
  header.version = kProtocolVersion;
  header.type = static_cast<uint8_t>(MessageType::Media);
  header.stream = static_cast<uint8_t>(StreamId::Audio);
  header.sequence = 9;
  header.pts_us = 42;
  uint8_t wire[kHeaderSize];
  SerializeHeader(header, wire);

  std::vector<uint8_t> message(wire, wire + kHeaderSize);
  message.insert(message.end(), payload.begin(), payload.end());
  for (size_t pos = 0; pos < message.size(); pos += 7) {
    size_t chunk = std::min<size_t>(7, message.size() - pos);
    REQUIRE(::send(fds[1], message.data() + pos, chunk, 0) == static_cast<ssize_t>(chunk));
    std::this_thread::sleep_for(std::chrono::milliseconds(1));
  }

  REQUIRE(collector.WaitFor([&] { return !collector.media.empty(); }));
  {
    std::lock_guard<std::mutex> lock(collector.mutex);
    REQUIRE(collector.media.size() == 1);
    REQUIRE(collector.media[0].first.sequence == 9);
    REQUIRE(collector.media[0].second == payload);
  }

  ::close(fds[1]);
  REQUIRE(collector.WaitFor([&] { return collector.disconnects >= 1; }));
  client.Stop();
}

TEST_CASE("StreamSocketClient disconnects on malformed input", "[stream_socket_client]") {
  int fds[2];
  REQUIRE(socketpair(AF_UNIX, SOCK_STREAM, 0, fds) == 0);

  Collector collector;
  StreamSocketClient client(fds[0], collector.MakeCallbacks());

  // Garbage: invalid version byte in an otherwise plausible header
  std::vector<uint8_t> garbage(kHeaderSize, 0xFF);
  REQUIRE(::send(fds[1], garbage.data(), garbage.size(), 0)
          == static_cast<ssize_t>(garbage.size()));

  REQUIRE(collector.WaitFor([&] { return collector.disconnects >= 1; }));
  {
    std::lock_guard<std::mutex> lock(collector.mutex);
    REQUIRE(collector.media.empty());
    REQUIRE(collector.hellos.empty());
  }
  ::close(fds[1]);
  client.Stop();
}

TEST_CASE("StreamSocketClient skips unknown message types", "[stream_socket_client]") {
  int fds[2];
  REQUIRE(socketpair(AF_UNIX, SOCK_STREAM, 0, fds) == 0);

  Collector collector;
  StreamSocketClient client(fds[0], collector.MakeCallbacks());

  // Unknown type 0x7f with a small payload, then a valid MEDIA
  Header header = {};
  header.length = kHeaderLengthBytes + 4;
  header.version = kProtocolVersion;
  header.type = 0x7f;
  uint8_t wire[kHeaderSize];
  SerializeHeader(header, wire);
  std::vector<uint8_t> message(wire, wire + kHeaderSize);
  message.insert(message.end(), {1, 2, 3, 4});

  header.type = static_cast<uint8_t>(MessageType::Media);
  header.length = kHeaderLengthBytes + 2;
  SerializeHeader(header, wire);
  message.insert(message.end(), wire, wire + kHeaderSize);
  message.insert(message.end(), {0xAA, 0xBB});

  REQUIRE(::send(fds[1], message.data(), message.size(), 0)
          == static_cast<ssize_t>(message.size()));

  REQUIRE(collector.WaitFor([&] { return !collector.media.empty(); }));
  {
    std::lock_guard<std::mutex> lock(collector.mutex);
    REQUIRE(collector.media.size() == 1);
    REQUIRE(collector.media[0].second == std::vector<uint8_t>({0xAA, 0xBB}));
  }
  ::close(fds[1]);
  client.Stop();
}

TEST_CASE("StreamSocketClient backs off when the producer rejects the connection",
          "[stream_socket_client]") {
  // A producer that accepts and immediately closes, as StreamSocket does for
  // a client failing the uid allow-list or the client limit
  ::unlink(kSockPath);
  int listener = ::socket(AF_UNIX, SOCK_STREAM, 0);
  REQUIRE(listener >= 0);
  sockaddr_un addr = {};
  addr.sun_family = AF_UNIX;
  strncpy(addr.sun_path, kSockPath, sizeof(addr.sun_path) - 1);
  REQUIRE(::bind(listener, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)) == 0);
  REQUIRE(::listen(listener, 8) == 0);

  std::atomic<int> accepts{0};
  std::atomic<bool> stop{false};
  std::thread rejecter([&] {
    while (!stop) {
      timeval tv = {0, 100000};
      fd_set fds;
      FD_ZERO(&fds);
      FD_SET(listener, &fds);
      if (::select(listener + 1, &fds, nullptr, nullptr, &tv) > 0) {
        int fd = ::accept(listener, nullptr, nullptr);
        if (fd >= 0) {
          ++accepts;
          ::close(fd);
        }
      }
    }
  });

  Collector collector;
  {
    StreamSocketClient client(std::string(kSockPath), collector.MakeCallbacks());
    std::this_thread::sleep_for(std::chrono::milliseconds(2500));
    client.Stop();
  }
  stop = true;
  rejecter.join();
  ::close(listener);
  ::unlink(kSockPath);

  // With the 1s backoff after every rejected connection: the initial attempt
  // plus two or three retries in 2.5s. Without backoff this is thousands.
  REQUIRE(accepts >= 2);
  REQUIRE(accepts <= 5);
  // A rejection is not a session, so no disconnect was reported
  std::lock_guard<std::mutex> lock(collector.mutex);
  REQUIRE(collector.disconnects == 0);
}

TEST_CASE("StreamSocketClient delivers EVENT frames", "[stream_socket_client]") {
  StreamSocket server(7, kSockPath);
  REQUIRE(server.Start());

  // A cached snapshot is replayed on connect, then a broadcast state change
  MonitorEvent snapshot;
  snapshot.code = kEventSnapshot;
  snapshot.state_id = 0;      snapshot.has_state_id = true;
  snapshot.state_name = "IDLE";
  snapshot.health_code = kEventCaptureFailed;  snapshot.has_health_code = true;
  snapshot.message = "Capture failed";
  server.SetSnapshotEvent(BuildEvent(snapshot));

  Collector collector;
  StreamSocketClient client(std::string(kSockPath), collector.MakeCallbacks());

  REQUIRE(collector.WaitFor([&] { return collector.events.size() >= 1; }));
  {
    std::lock_guard<std::mutex> lock(collector.mutex);
    const Header &header = collector.events[0].first;
    const MonitorEvent &event = collector.events[0].second;
    REQUIRE(header.stream == static_cast<uint8_t>(StreamId::Monitor));
    REQUIRE(event.code == kEventSnapshot);
    REQUIRE(event.state_name == "IDLE");
    REQUIRE(event.has_health_code);
    REQUIRE(event.health_code == kEventCaptureFailed);
    REQUIRE(event.message == "Capture failed");
  }

  std::this_thread::sleep_for(std::chrono::milliseconds(100));
  MonitorEvent change;
  change.code = kEventStateChanged;
  change.state_id = 2;       change.has_state_id = true;
  change.prev_state_id = 0;  change.has_prev_state_id = true;
  change.state_name = "ALARM";
  change.wall_clock_us = 1718355103501000ULL;  change.has_wall_clock = true;
  server.SendMonitorEvent(BuildEvent(change));

  REQUIRE(collector.WaitFor([&] { return collector.events.size() >= 2; }));
  {
    std::lock_guard<std::mutex> lock(collector.mutex);
    const Header &header = collector.events[1].first;
    const MonitorEvent &event = collector.events[1].second;
    REQUIRE(header.sequence == 0);  // first event produced by this monitor
    REQUIRE(event.code == kEventStateChanged);
    REQUIRE(event.state_id == 2);
    REQUIRE(event.prev_state_id == 0);
    REQUIRE(event.state_name == "ALARM");
    REQUIRE(event.wall_clock_us == 1718355103501000ULL);
  }

  client.Stop();
  server.Stop();
}

TEST_CASE("StreamSocketClient refuses a path too long for a unix socket address", "[stream_socket_client]") {
  Collector collector;
  StreamSocketClient client("/tmp/" + std::string(200, 'x') + ".sock", collector.MakeCallbacks());
  // The reader thread gives up instead of connecting to a truncated path
  std::this_thread::sleep_for(std::chrono::milliseconds(200));
  REQUIRE_FALSE(client.IsConnected());
  client.Stop();
  REQUIRE(collector.disconnects == 0);
}