#include "cmCTestJobServerPosix.h"

#include <algorithm>
#include <cassert>
#include <cstring>
#include <string>
#include <utility>

#include <cm/optional>

#include <fcntl.h>

#include "cmStringAlgorithms.h"
#include "cmSystemTools.h"

uv_stream_t* cmCTestJobServerPosix::_get_writer()
{
  switch (connection_type) {
    case ConnectionType::Pipe:
      return reinterpret_cast<uv_stream_t*>(&write_pipe);
    case ConnectionType::Fifo:
      return reinterpret_cast<uv_stream_t*>(&fifo);
    default:
      return nullptr;
  }
}

uv_stream_t* cmCTestJobServerPosix::_get_reader()
{
  switch (connection_type) {
    case ConnectionType::Pipe:
      return reinterpret_cast<uv_stream_t*>(&read_pipe);
    case ConnectionType::Fifo:
      return reinterpret_cast<uv_stream_t*>(&fifo);
    default:
      return nullptr;
  }
}

void cmCTestJobServerPosix::_on_read(uv_stream_t* stream, ssize_t nread,
                                     const uv_buf_t* buf)
{
  cmCTestJobServerPosix* self =
    reinterpret_cast<cmCTestJobServerPosix*>(stream->data);

  if (nread < 0) {
    if (self->FatalCallback.has_value()) {
      cmSystemTools::Error(
        cmStrCat("Failed to read jobserver tokens: ", uv_strerror(nread)));
      self->FatalCallback.value()();
    }
    return;
  }

  if (nread > 0) {
    self->tokens.insert(self->tokens.end(), buf->base, buf->base + nread);
    self->_process_queue();
  }

  delete buf->base;
}

// Marks tokens as used, this is solely bookkeeping and doesn't actually read
void cmCTestJobServerPosix::_acquire(size_t slots)
{
  if (slots == 0) {
    return;
  }

  // Prefer to use the free token, if available
  if (this->freeTokenAvailable) {
    this->freeTokenAvailable = false;
    slots--;
  }

  this->usedTokens += slots;
}

// Starts as many jobs as possible, given the number of tokens available
//
// If we exhaust the queue, we release any unused tokens
// If we were unable to start the next job, and no jobs are running, we might
// have entered a deadlock. So we start the next job with as many tokens as we
// have available
void cmCTestJobServerPosix::_process_queue()
{
  size_t slots = freeTokenAvailable + this->tokens.size() - this->usedTokens;
  size_t started = 0;
  while (slots > 0 && !queue.empty()) {
    size_t slotsNeeded = std::min(queue.front().slots, maxSlots);
    if (slotsNeeded > slots) {
      break;
    }

    cmCTestJobServerClient::Task task = std::move(queue.front());
    queue.pop_front();
    slots -= task.slots;
    _acquire(task.slots);
    task.run(task.slots);

    started++;
  }

  if (queue.empty()) {
    // Be a friend and release any tokens we didn't use
    _write(this->tokens.size() - this->usedTokens, true);
  } else if (started == 0 && this->freeTokenAvailable &&
             this->usedTokens == 0) {
    // If no jobs were started, and we're not running any jobs, then we start
    // the next job with as many tokens as we have available
    cmCTestJobServerClient::Task task = std::move(queue.front());
    queue.pop_front();
    _acquire(slots);
    task.run(slots);
  }
}

// Synchronously write the tokens to the pipe, blocking here is convenient
// for implementing `Close()`. There shouldn't be much of a performance
// impact because there are always other writers on the pipe, and the number
// of tokens is small.
//
// If `retry` is false then we only attempt to write once and return
void cmCTestJobServerPosix::_write(size_t n, bool retry)
{
  if (n == 0) {
    return;
  }

  // Takes n tokens from the queue and copies them into a buffer
  char buf[n];
  std::memcpy(buf, this->tokens.data(), n);
  this->tokens.erase(this->tokens.begin(), this->tokens.begin() + n);

  uv_stream_t* writer = _get_writer();
  uv_buf_t uvbuf = uv_buf_init(buf, n);
  size_t total = 0;
  do {
    int written = uv_try_write(writer, &uvbuf, 1);
    if (written < 0 && written != UV_EAGAIN) {
      cmSystemTools::Error(
        cmStrCat("Failed to write jobserver tokens: ", uv_strerror(written)));
      if (FatalCallback.has_value()) {
        FatalCallback.value()();
      }
      return;
    }

    total += written;
    uvbuf.base += written;
    uvbuf.len -= written;
  } while (retry && total < n);
}

bool cmCTestJobServerPosix::Connect(uv_loop_t* loop, int rfd, int wfd)
{
  this->Loop = loop;
  this->connection_type = ConnectionType::Pipe;

  if (uv_pipe_init(loop, &read_pipe, 0) != 0 ||
      uv_pipe_init(loop, &write_pipe, 0) != 0) {
    return false;
  }

  if (uv_pipe_open(&read_pipe, rfd) != 0 ||
      uv_pipe_open(&write_pipe, wfd) != 0) {
    return false;
  }

  // Verify that the read side is readable, and the write side is writable
  if (!uv_is_readable(reinterpret_cast<uv_stream_t*>(&read_pipe)) ||
      !uv_is_writable(reinterpret_cast<uv_stream_t*>(&write_pipe))) {
    return false;
  }

  read_pipe.data = this;
  write_pipe.data = this;

  uv_read_start(
    (uv_stream_t*)&read_pipe,
    [](uv_handle_t* /*handle*/, size_t suggested_size, uv_buf_t* buf) {
      buf->base = new char[suggested_size];
      buf->len = suggested_size;
    },
    _on_read);

  return true;
}

bool cmCTestJobServerPosix::Connect(uv_loop_t* loop, const char* path)
{
  this->Loop = loop;
  this->connection_type = ConnectionType::Fifo;

  if (uv_pipe_init(loop, &fifo, 0) != 0) {
    return false;
  }

  int fd = open(path, O_RDWR);
  if (fd == -1) {
    return false;
  }

  if (uv_pipe_open(&fifo, fd) != 0) {
    return false;
  }

  uv_read_start(
    (uv_stream_t*)&fifo,
    [](uv_handle_t* /*handle*/, size_t suggested_size, uv_buf_t* buf) {
      buf->base = new char[suggested_size];
      buf->len = suggested_size;
    },
    _on_read);

  fifo.data = this;
  return true;
}

bool cmCTestJobServerPosix::Connect(uv_loop_t* loop)
{
  // --jobserver-auth= for gnu make versions >= 4.2
  // --jobserver-fds= for gnu make versions < 4.2
  // -J for bsd make
  const std::vector<cm::string_view> prefixes = { "--jobserver-auth=",
                                                  "--jobserver-fds=", "-J" };

  cm::optional<std::string> maxJobs = cm::nullopt;
  cm::optional<std::string> auth = cm::nullopt;

  std::string makeflags;
  if (!cmSystemTools::GetEnv("MAKEFLAGS", makeflags)) {
    return false;
  }

  std::vector<std::string> args;
  cmSystemTools::ParseUnixCommandLine(makeflags.c_str(), args);
  for (const std::string& arg : args) {
    cm::string_view arg_view(arg);

    if (cmHasLiteralPrefix(arg_view, "-j")) {
      auth = arg_view.substr(cmStrLen("-j"));
    } else {
      for (const cm::string_view& prefix : prefixes) {
        if (cmHasPrefix(arg_view, prefix)) {
          auth = cmTrimWhitespace(arg_view.substr(prefix.length()));
          break;
        }
      }
    }
  }

  if (maxJobs) {
    try {
      this->maxSlots = std::min(this->maxSlots, std::stoul(*maxJobs));
    } catch (...) {
    }
  }

  if (!auth) {
    return false;
  }

  // fifo:PATH
  if (cmHasLiteralPrefix(*auth, "fifo:")) {
    return Connect(loop, auth->substr(cmStrLen("fifo:")).c_str());
  }

  // reader,writer
  size_t reader;
  size_t writer;
  try {
    size_t comma = auth->find(',');
    if (comma == std::string::npos) {
      return false;
    }

    reader = std::stoul(auth->substr(0, comma));
    writer = std::stoul(auth->substr(comma + 1));
  } catch (...) {
    return false;
  }

  return Connect(loop, reader, writer);
}

void cmCTestJobServerPosix::Enqueue(std::function<void(size_t)> task,
                                    size_t slots)
{
  this->queue.push_back({ task, slots });
}

void cmCTestJobServerPosix::Release(size_t n)
{
  if (n == 0) {
    return;
  }

  if (n > this->usedTokens) {
    assert(n == this->usedTokens + 1);
    this->freeTokenAvailable = true;
    n--;
  }
  this->usedTokens -= n;

  _write(n, true);
  _process_queue();
}

void cmCTestJobServerPosix::Close(bool force)
{
  uv_read_stop(_get_reader());
  _write(this->tokens.size(), !force);

  switch (this->connection_type) {
    case ConnectionType::Fifo:
      uv_close(reinterpret_cast<uv_handle_t*>(&fifo),
               [](uv_handle_t* /*handle*/) {});
      break;
    case ConnectionType::Pipe:
      uv_close(reinterpret_cast<uv_handle_t*>(&read_pipe),
               [](uv_handle_t* /*handle*/) {});

      uv_close(reinterpret_cast<uv_handle_t*>(&write_pipe),
               [](uv_handle_t* /*handle*/) {});
      break;
  }
}

void cmCTestJobServerPosix::SetFatalCallback(std::function<void()> f)
{
  this->FatalCallback = std::move(f);
}
