// This source code is released into the public domain.
export module nihil.posix:spawn;

// spawn(): fork and execute a child process.

import nihil.std;
import nihil.util;
import :argv;
import :executor;
import :fd;
import :open;
import :process;
import :unistd;

namespace nihil {

// fd_{read,write}_pipe: create a pipe with one end in the child and the other in the parent.
// The child's side will be dup2()'d to the provided fd number.
// The parent side fd can be retrieved via parent_fd().
//
// fd_read_pipe() puts the reading side in the child, while fd_write_pipe() puts the writing
// side in the child.

struct fd_pipe_base
{
	fd_pipe_base(int fdno, fd &&child_fd, fd &&parent_fd)
		: m_fdno(fdno)
		, m_child_fd(std::move(child_fd))
		, m_parent_fd(std::move(parent_fd))
	{
	}

	auto run_in_child(this fd_pipe_base &self, process &) -> void
	{
		auto err = raw_dup(self.m_child_fd, self.m_fdno);
		if (!err) {
			std::println("dup: {}", err.error());
			std::quick_exit(1);
		}

		/*
		 * We don't care about errors from close() since the fd
		 * is still closed.
		 */
		std::ignore = self.m_parent_fd.close();
		std::ignore = self.m_child_fd.close();
	}

	auto run_in_parent(this fd_pipe_base &self, process &) -> void
	{
		std::ignore = self.m_child_fd.close();
	}

	[[nodiscard]] auto parent_fd(this fd_pipe_base &self) -> fd &
	{
		return self.m_parent_fd;
	}

private:
	int m_fdno;
	fd m_child_fd;
	fd m_parent_fd;
};

export struct fd_read_pipe final : fd_pipe_base
{
	fd_read_pipe(int fdno, fd &&read_fd, fd &&write_fd)
		: fd_pipe_base(fdno, std::move(read_fd), std::move(write_fd))
	{
	}
};

export struct fd_write_pipe final : fd_pipe_base
{
	fd_write_pipe(int fdno, fd &&read_fd, fd &&write_fd)
		: fd_pipe_base(fdno, std::move(write_fd), std::move(read_fd))
	{
	}
};

export [[nodiscard]] auto make_fd_read_pipe(int fdno) -> std::expected<fd_read_pipe, error>
{
	auto fds = co_await pipe();
	co_return fd_read_pipe(fdno, std::move(fds.first), std::move(fds.second));
}

export [[nodiscard]] auto make_fd_write_pipe(int fdno) -> std::expected<fd_write_pipe, error>
{
	auto fds = co_await pipe();
	co_return fd_write_pipe(fdno, std::move(fds.first), std::move(fds.second));
}

/*
 * fd_file: open a file and provide it to the child as a file descriptor.
 * open_flags and open_mode are as for ::open().
 */
export struct fd_file final
{
	fd_file(int fdno, fd &&file_fd)
		: m_fdno(fdno)
		, m_file_fd(std::move(file_fd))
	{
	}

	auto run_in_parent(this fd_file &self, process &) -> void
	{
		std::ignore = self.m_file_fd.close();
	}

	auto run_in_child(this fd_file &self, process &) -> void
	{
		auto err = raw_dup(self.m_file_fd, self.m_fdno);
		if (!err) {
			std::print("dup: {}\n", err.error());
			std::quick_exit(1);
		}

		std::ignore = self.m_file_fd.close();
	}

private:
	int m_fdno;
	fd m_file_fd;
};

export [[nodiscard]] auto
make_fd_file(int fdno, std::filesystem::path const &file, open_flags flags, int mode = 0777)
	-> std::expected<fd_file, error>
{
	auto fd = co_await open(file, flags, mode);
	co_return fd_file(fdno, std::move(fd));
}

/*
 * Shorthand for fd_file with /dev/null as the file.
 */

export [[nodiscard]] inline auto stdin_devnull() -> std::expected<fd_file, error>
{
	return make_fd_file(stdin_fileno, "/dev/null", open_read);
}

export [[nodiscard]] inline auto stdout_devnull() -> std::expected<fd_file, error>
{
	return make_fd_file(stdout_fileno, "/dev/null", open_write);
}

export [[nodiscard]] inline auto stderr_devnull() -> std::expected<fd_file, error>
{
	return make_fd_file(stderr_fileno, "/dev/null", open_write);
}

/*
 * Capture the output of a pipe in the parent and read it into an
 * output iterator.
 */
export template <std::output_iterator<char> Iterator>
struct fd_capture final
{
	fd_capture(fd_write_pipe &&pipe, Iterator it)
		: m_pipe(std::move(pipe))
		, m_iterator(std::move(it))
	{
	}

	auto run_in_child(this fd_capture &self, process &p) -> void
	{
		self.m_pipe.run_in_child(p);
	}

	auto run_in_parent(this fd_capture &self, process &p) -> void
	{
		self.m_pipe.run_in_parent(p);

		auto constexpr bufsize = std::size_t{1024};
		auto buffer = std::array<char, bufsize>();
		auto &fd = self.m_pipe.parent_fd();
		for (;;) {
			auto ret = read(fd, buffer);
			if (!ret || ret->empty())
				break;

			std::ranges::copy(*ret, self.m_iterator);
		}

		// We probably want to handle errors here somehow,
		// but it's not clear what would be useful behaviour.
	}

private:
	fd_write_pipe m_pipe;
	Iterator m_iterator;
};

export [[nodiscard]] auto make_capture(int fdno, std::output_iterator<char> auto &&it)
	-> std::expected<fd_capture<decltype(it)>, error>
{
	auto pipe = co_await make_fd_write_pipe(fdno);
	co_return fd_capture(std::move(pipe), std::forward<decltype(it)>(it));
}

export [[nodiscard]] auto make_capture(int fdno, std::string &str)
	-> std::expected<fd_capture<decltype(std::back_inserter(str))>, error>
{
	auto pipe = co_await make_fd_write_pipe(fdno);
	co_return fd_capture(std::move(pipe), std::back_inserter(str));
}

/*
 * Spawn a new process with the given arguments and return a struct process.
 * Throws exec_error() on failure.
 */
export [[nodiscard]] auto
spawn(executor auto &&executor, auto &&...actions) -> std::expected<process, error>
{
	auto const pid = co_await fork();
	auto proc = process(pid);

	if (pid == 0) {
		// We are in the child.
		(actions.run_in_child(proc), ...);
		std::ignore = std::move(proc).release();

		auto err = executor.exec();
		std::println("{}", err.error());
		std::quick_exit(1);
	}

	(actions.run_in_parent(proc), ...);

	co_return proc;
}

} // namespace nihil