#include "spawner.hh"

#include "image_processor.hh"
#include "io.hh"
#include "unique_fd.hh"

#include <cassert>
#include <csignal>
#include <expected>
#include <linux/prctl.h>
#include <memory>
#include <sys/prctl.h>
#include <sys/syscall.h>
#include <unistd.h>
#include <utility>

namespace {

int pidfd_getfd(int pidfd, int targetfd, unsigned int flags) {
  return static_cast<int>(syscall(SYS_pidfd_getfd, pidfd, targetfd, flags));
}

// NOLINTNEXTLINE(misc-include-cleaner)
int pidfd_open(pid_t pid, unsigned int flags) {
  return static_cast<int>(syscall(SYS_pidfd_open, pid, flags));
}

class ProcessImpl : public Process {
 public:
  ProcessImpl(
      pid_t pid,
      std::pair<std::unique_ptr<io::Reader>, std::unique_ptr<io::Writer>> pipe)
      : pid_(pid), pipe_(std::move(pipe)) {
    assert(pid_);
  }

  [[nodiscard]]
  pid_t pid() const { return pid_; }

  [[nodiscard]]
  io::Reader& reader() const override { return *pipe_.first; }

  [[nodiscard]]
  io::Writer& writer() const override { return *pipe_.second; }

  void kill() {
    if (pid_ != 0) {
      ::kill(pid_, SIGTERM); // NOLINT(misc-include-cleaner)
      pid_ = 0;
    }
  }

 private:
  pid_t pid_;
  std::pair<std::unique_ptr<io::Reader>, std::unique_ptr<io::Writer>> pipe_;
};

struct Request {
  Spawner::Exec exec;
  int reader;
  int writer;
};

struct Response {
  pid_t pid;
};

void spawner_runner(unique_fd parent, std::unique_ptr<io::Reader> reader,
                    std::unique_ptr<io::Writer> writer) {
  while (true) {
    Request request;
    {
      auto ret = reader->repeat_read(&request, sizeof(request));
      if (!ret.has_value() || ret.value() != sizeof(request)) {
        break;
      }
    }

    // Need to not return a response before child has duped fds.
    // So use a short-lived pipe to sync child start.
    auto sync_pipe = io::pipe();
    pid_t child_pid = 0;

    if (sync_pipe.has_value()) {
      child_pid = fork();
      if (child_pid == 0) {
        // Child process
        sync_pipe->first.reset();
        reader.reset();
        writer.reset();

        unique_fd child_reader_fd{pidfd_getfd(parent.get(), request.reader, 0)};
        unique_fd child_writer_fd{pidfd_getfd(parent.get(), request.writer, 0)};

        parent.reset();

        if (child_reader_fd && child_writer_fd) {
          char c = 1;
          if (sync_pipe->second->write(&c, sizeof(c)).has_value()) {
            auto child_reader = io::reader_from_raw(child_reader_fd.release());
            auto child_writer = io::writer_from_raw(child_writer_fd.release());

            switch (request.exec) {
              case Spawner::Exec::kImageProcessor:
                _exit(image_processor::run(std::move(child_reader),
                                           std::move(child_writer)));
            }
          }
        }

        _exit(1);
        // _exit obviously never returns but to help tools and compilers…
        return;
      }

      // Parent process
      if (child_pid == -1) {
        // fork() failed, we use zero as error value.
        child_pid = 0;
      } else {
        sync_pipe->second.reset();

        char c = 1;
        std::ignore = sync_pipe->first->read(&c, sizeof(c));
      }
    }

    Response response{child_pid};
    {
      auto ret = writer->repeat_write(&response, sizeof(response));
      if (!ret.has_value() || ret.value() != sizeof(response)) {
        break;
      }
    }
  }
}

class SpawnerImpl : public Spawner {
 public:
  explicit SpawnerImpl(std::unique_ptr<ProcessImpl> process)
      : process_(std::move(process)) {}

  ~SpawnerImpl() override {
    if (process_) {
      process_->kill();
    }
  }

  [[nodiscard]]
  std::expected<std::unique_ptr<Process>, Error> run(Exec exec) override {
    if (!process_) {
      return std::unexpected(Error::kError);
    }

    auto reader_pipe = io::pipe();
    auto writer_pipe = io::pipe();
    if (!reader_pipe.has_value() || !writer_pipe.has_value()) {
      return std::unexpected(Error::kError);
    }

    Request req{
      .exec = exec,
      .reader = writer_pipe->first->raw_fd(),
      .writer = reader_pipe->second->raw_fd(),
    };
    {
      auto ret = process_->writer().repeat_write(&req, sizeof(req));
      if (!ret.has_value() || ret.value() != sizeof(req)) {
        process_->kill();
        process_.reset();
        return std::unexpected(Error::kError);
      }
    }

    Response resp;
    {
      auto ret = process_->reader().repeat_read(&resp, sizeof(resp));
      if (!ret.has_value() || ret.value() != sizeof(resp)) {
        process_->kill();
        process_.reset();
        return std::unexpected(Error::kError);
      }
    }

    if (resp.pid == 0)
      return std::unexpected(Error::kError);

    return std::make_unique<ProcessImpl>(
        resp.pid, std::make_pair(std::move(reader_pipe->first),
                                 std::move(writer_pipe->second)));
  }

 private:
  std::unique_ptr<ProcessImpl> process_;
};

}  // namespace

std::expected<std::unique_ptr<Spawner>, Spawner::Error> Spawner::create() {
  auto reader_pipe = io::pipe();
  auto writer_pipe = io::pipe();
  if (!reader_pipe.has_value() || !writer_pipe.has_value()) {
    return std::unexpected(Error::kError);
  }

  auto pid = getpid();
  unique_fd pidfd{pidfd_open(pid, 0)};
  if (!pidfd) {
    return std::unexpected(Error::kError);
  }

  // Needed for pidfd_getfd to work in child.
  if (prctl(PR_SET_PTRACER, pid) != 0) {
    return std::unexpected(Error::kError);
  }

  pid = fork();
  if (pid == 0) {
    // Child process
    reader_pipe->first.reset();
    writer_pipe->second.reset();

    spawner_runner(std::move(pidfd), std::move(writer_pipe->first),
                   std::move(reader_pipe->second));
    _exit(0);
  }

  // Parent process
  if (pid == -1) {
    return std::unexpected(Error::kError);
  }

  return std::make_unique<SpawnerImpl>(std::make_unique<ProcessImpl>(
      pid, std::make_pair(std::move(reader_pipe->first),
                          std::move(writer_pipe->second))));
}