#include "cvmmap_streamer/protocol/rtmp_output.hpp"

#include "cvmmap_streamer/protocol/rtmp_publisher.hpp"

extern "C" {
#include <libavcodec/avcodec.h>
#include <libavformat/avformat.h>
#include <libavutil/avutil.h>
}

#include <algorithm>
#include <cerrno>
#include <cstddef>
#include <cstdint>
#include <cstring>
#include <span>
#include <string>
#include <string_view>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>
#include <utility>
#include <vector>

#include <fcntl.h>
#include <spdlog/spdlog.h>

namespace cvmmap_streamer::protocol {

namespace {

constexpr std::uint64_t kNanosPerSecond = 1'000'000'000ull;

[[nodiscard]]
std::string av_error_string(int error_code) {
	char buffer[AV_ERROR_MAX_STRING_SIZE]{};
	av_strerror(error_code, buffer, sizeof(buffer));
	return std::string(buffer);
}

[[nodiscard]]
AVCodecID to_avcodec_id(CodecType codec) {
	return codec == CodecType::H265 ? AV_CODEC_ID_HEVC : AV_CODEC_ID_H264;
}

[[nodiscard]]
const char *ffmpeg_input_format(CodecType codec) {
	return codec == CodecType::H265 ? "hevc" : "h264";
}

[[nodiscard]]
AVRational stream_time_base(const encode::EncodedStreamInfo &stream_info) {
	return AVRational{
		static_cast<int>(stream_info.time_base_num),
		static_cast<int>(stream_info.time_base_den),
	};
}

[[nodiscard]]
Status copy_decoder_config(
	AVCodecParameters *codecpar,
	std::span<const std::uint8_t> decoder_config) {
	if (codecpar == nullptr) {
		return unexpected_error(ERR_INVALID_ARGUMENT, "RTMP stream codec parameters are null");
	}
	if (decoder_config.empty()) {
		return {};
	}

	codecpar->extradata = static_cast<std::uint8_t *>(av_mallocz(decoder_config.size() + AV_INPUT_BUFFER_PADDING_SIZE));
	if (codecpar->extradata == nullptr) {
		return unexpected_error(ERR_ALLOCATION_FAILED, "failed to allocate RTMP codec extradata");
	}
	std::memcpy(codecpar->extradata, decoder_config.data(), decoder_config.size());
	codecpar->extradata_size = static_cast<int>(decoder_config.size());
	return {};
}

void append_start_code(std::vector<std::uint8_t> &output) {
	output.push_back(0x00);
	output.push_back(0x00);
	output.push_back(0x00);
	output.push_back(0x01);
}

[[nodiscard]]
Result<std::uint16_t> read_be16(std::span<const std::uint8_t> bytes, std::size_t offset) {
	if (offset + 2 > bytes.size()) {
		return unexpected_error(ERR_PROTOCOL, "decoder config truncated");
	}
	return static_cast<std::uint16_t>((static_cast<std::uint16_t>(bytes[offset]) << 8) | bytes[offset + 1]);
}

[[nodiscard]]
Result<std::vector<std::uint8_t>> avcc_to_annexb(std::span<const std::uint8_t> decoder_config) {
	if (decoder_config.size() < 7 || decoder_config[0] != 1) {
		return unexpected_error(ERR_PROTOCOL, "invalid AVC decoder config");
	}

	std::vector<std::uint8_t> annexb{};
	std::size_t offset = 5;
	const auto sps_count = static_cast<std::size_t>(decoder_config[offset++] & 0x1fu);
	for (std::size_t i = 0; i < sps_count; ++i) {
		auto size = read_be16(decoder_config, offset);
		if (!size) {
			return std::unexpected(size.error());
		}
		offset += 2;
		if (offset + *size > decoder_config.size()) {
			return unexpected_error(ERR_PROTOCOL, "invalid AVC decoder config payload");
		}
		append_start_code(annexb);
		annexb.insert(annexb.end(), decoder_config.begin() + static_cast<std::ptrdiff_t>(offset), decoder_config.begin() + static_cast<std::ptrdiff_t>(offset + *size));
		offset += *size;
	}

	if (offset >= decoder_config.size()) {
		return unexpected_error(ERR_PROTOCOL, "invalid AVC decoder config: missing PPS count");
	}
	const auto pps_count = static_cast<std::size_t>(decoder_config[offset++]);
	for (std::size_t i = 0; i < pps_count; ++i) {
		auto size = read_be16(decoder_config, offset);
		if (!size) {
			return std::unexpected(size.error());
		}
		offset += 2;
		if (offset + *size > decoder_config.size()) {
			return unexpected_error(ERR_PROTOCOL, "invalid AVC decoder config payload");
		}
		append_start_code(annexb);
		annexb.insert(annexb.end(), decoder_config.begin() + static_cast<std::ptrdiff_t>(offset), decoder_config.begin() + static_cast<std::ptrdiff_t>(offset + *size));
		offset += *size;
	}

	return annexb;
}

[[nodiscard]]
Result<std::vector<std::uint8_t>> hvcc_to_annexb(std::span<const std::uint8_t> decoder_config) {
	if (decoder_config.size() < 23 || decoder_config[0] != 1) {
		return unexpected_error(ERR_PROTOCOL, "invalid HEVC decoder config");
	}

	std::vector<std::uint8_t> annexb{};
	std::size_t offset = 22;
	const auto array_count = static_cast<std::size_t>(decoder_config[offset++]);
	for (std::size_t array_index = 0; array_index < array_count; ++array_index) {
		if (offset + 3 > decoder_config.size()) {
			return unexpected_error(ERR_PROTOCOL, "invalid HEVC decoder config arrays");
		}
		offset += 1;
		auto nal_count = read_be16(decoder_config, offset);
		if (!nal_count) {
			return std::unexpected(nal_count.error());
		}
		offset += 2;

		for (std::size_t nal_index = 0; nal_index < *nal_count; ++nal_index) {
			auto nal_size = read_be16(decoder_config, offset);
			if (!nal_size) {
				return std::unexpected(nal_size.error());
			}
			offset += 2;
			if (offset + *nal_size > decoder_config.size()) {
				return unexpected_error(ERR_PROTOCOL, "invalid HEVC decoder config payload");
			}
			append_start_code(annexb);
			annexb.insert(annexb.end(), decoder_config.begin() + static_cast<std::ptrdiff_t>(offset), decoder_config.begin() + static_cast<std::ptrdiff_t>(offset + *nal_size));
			offset += *nal_size;
		}
	}

	return annexb;
}

[[nodiscard]]
bool looks_like_annexb(std::span<const std::uint8_t> bytes) {
	if (bytes.size() >= 4 &&
		bytes[0] == 0x00 &&
		bytes[1] == 0x00 &&
		bytes[2] == 0x00 &&
		bytes[3] == 0x01) {
		return true;
	}
	return bytes.size() >= 3 &&
		   bytes[0] == 0x00 &&
		   bytes[1] == 0x00 &&
		   bytes[2] == 0x01;
}

[[nodiscard]]
Result<std::vector<std::uint8_t>> decoder_config_to_annexb(
	CodecType codec,
	std::span<const std::uint8_t> decoder_config) {
	if (decoder_config.empty()) {
		return unexpected_error(ERR_PROTOCOL, "decoder config is required");
	}
	if (looks_like_annexb(decoder_config)) {
		return std::vector<std::uint8_t>(decoder_config.begin(), decoder_config.end());
	}
	if (codec == CodecType::H265) {
		return hvcc_to_annexb(decoder_config);
	}
	return avcc_to_annexb(decoder_config);
}

[[nodiscard]]
Status write_all(int fd, std::span<const std::uint8_t> bytes) {
	std::size_t written{0};
	while (written < bytes.size()) {
		const auto result = ::write(fd, bytes.data() + written, bytes.size() - written);
		if (result < 0) {
			if (errno == EINTR) {
				continue;
			}
			return unexpected_error(ERR_IO, std::strerror(errno));
		}
		if (result == 0) {
			return unexpected_error(ERR_IO, "short write to ffmpeg process stdin");
		}
		written += static_cast<std::size_t>(result);
	}
	return {};
}

class LegacyCustomRtmpOutput {
public:
	explicit LegacyCustomRtmpOutput(RtmpPublisher &&publisher)
		: publisher_(std::move(publisher)) {}

	[[nodiscard]]
	std::string_view backend_name() const {
		return "legacy_custom";
	}

	[[nodiscard]]
	Status publish_access_unit(const encode::EncodedAccessUnit &access_unit) {
		auto publish = publisher_.publish_access_unit(access_unit.annexb_bytes, access_unit.stream_pts_ns);
		if (!publish) {
			return unexpected_error(ERR_PROTOCOL, publish.error());
		}
		return {};
	}

	void log_metrics() const {
		publisher_.log_metrics();
	}

private:
	RtmpPublisher publisher_{};
};

class LibavformatRtmpOutput {
public:
	struct Session {
		std::string url{};
		AVFormatContext *format_context{nullptr};
		AVStream *video_stream{nullptr};
	};

	LibavformatRtmpOutput() = default;
	LibavformatRtmpOutput(const LibavformatRtmpOutput &) = delete;
	LibavformatRtmpOutput &operator=(const LibavformatRtmpOutput &) = delete;
	LibavformatRtmpOutput(LibavformatRtmpOutput &&) noexcept = default;
	LibavformatRtmpOutput &operator=(LibavformatRtmpOutput &&) noexcept = default;

	~LibavformatRtmpOutput() {
		for (auto &session : sessions_) {
			close_session(session);
		}
	}

	[[nodiscard]]
	static Result<RtmpOutput> create(
		const RuntimeConfig &config,
		const encode::EncodedStreamInfo &stream_info) {
		if (stream_info.decoder_config.empty()) {
			return unexpected_error(ERR_PROTOCOL, "libavformat RTMP requires encoder decoder_config/extradata");
		}

		avformat_network_init();

		LibavformatRtmpOutput output{};
		output.codec_ = stream_info.codec;

		for (const auto &url : config.outputs.rtmp.urls) {
			auto session = create_session(url, stream_info);
			if (!session) {
				return std::unexpected(session.error());
			}
			output.sessions_.push_back(std::move(*session));
		}

		return pro::make_proxy<RtmpOutputFacade, LibavformatRtmpOutput>(std::move(output));
	}

	[[nodiscard]]
	std::string_view backend_name() const {
		return "libavformat";
	}

	[[nodiscard]]
	Status publish_access_unit(const encode::EncodedAccessUnit &access_unit) {
		for (auto &session : sessions_) {
			auto packet = av_packet_alloc();
			if (packet == nullptr) {
				return unexpected_error(ERR_ALLOCATION_FAILED, "failed to allocate RTMP AVPacket");
			}
			const auto packet_result = av_new_packet(packet, static_cast<int>(access_unit.annexb_bytes.size()));
			if (packet_result < 0) {
				av_packet_free(&packet);
				return unexpected_error(ERR_ALLOCATION_FAILED, "failed to allocate RTMP packet payload: " + av_error_string(packet_result));
			}

			std::memcpy(packet->data, access_unit.annexb_bytes.data(), access_unit.annexb_bytes.size());
			packet->stream_index = session.video_stream->index;
			packet->flags        = access_unit.keyframe ? AV_PKT_FLAG_KEY : 0;
			packet->pts          = av_rescale_q(
				static_cast<std::int64_t>(access_unit.stream_pts_ns),
				AVRational{1, static_cast<int>(kNanosPerSecond)},
				session.video_stream->time_base);
			packet->dts = packet->pts;

			const auto write_result = av_interleaved_write_frame(session.format_context, packet);
			av_packet_free(&packet);
			if (write_result < 0) {
				return unexpected_error(
					ERR_NETWORK,
					"libavformat RTMP write failed for '" + session.url + "': " + av_error_string(write_result));
			}

			video_messages_ += 1;
			bytes_sent_ += access_unit.annexb_bytes.size();
		}

		access_units_ += 1;
		return {};
	}

	void log_metrics() const {
		spdlog::info(
			"RTMP_OUTPUT_METRICS backend={} codec={} urls={} access_units={} video_messages={} bytes_sent={}",
			backend_name(),
			to_string(codec_),
			sessions_.size(),
			access_units_,
			video_messages_,
			bytes_sent_);
	}

private:
	[[nodiscard]]
	static Result<Session> create_session(
		const std::string &url,
		const encode::EncodedStreamInfo &stream_info) {
		Session session{};
		session.url = url;

		const auto alloc_result = avformat_alloc_output_context2(&session.format_context, nullptr, "flv", url.c_str());
		if (alloc_result < 0 || session.format_context == nullptr) {
			return unexpected_error(
				ERR_ALLOCATION_FAILED,
				"failed to allocate RTMP output context for '" + url + "': " + av_error_string(alloc_result));
		}

		session.format_context->flags |= AVFMT_FLAG_FLUSH_PACKETS;

		session.video_stream = avformat_new_stream(session.format_context, nullptr);
		if (session.video_stream == nullptr) {
			close_session(session);
			return unexpected_error(ERR_ALLOCATION_FAILED, "failed to allocate RTMP output stream for '" + url + "'");
		}

		session.video_stream->time_base      = stream_time_base(stream_info);
		session.video_stream->avg_frame_rate = AVRational{
			static_cast<int>(stream_info.frame_rate_num),
			static_cast<int>(stream_info.frame_rate_den),
		};

		auto *codecpar      = session.video_stream->codecpar;
		codecpar->codec_type = AVMEDIA_TYPE_VIDEO;
		codecpar->codec_id   = to_avcodec_id(stream_info.codec);
		codecpar->width      = static_cast<int>(stream_info.width);
		codecpar->height     = static_cast<int>(stream_info.height);

		auto extradata_copy = copy_decoder_config(codecpar, stream_info.decoder_config);
		if (!extradata_copy) {
			close_session(session);
			return std::unexpected(extradata_copy.error());
		}

		if (!(session.format_context->oformat->flags & AVFMT_NOFILE)) {
			const auto open_result = avio_open2(&session.format_context->pb, url.c_str(), AVIO_FLAG_WRITE, nullptr, nullptr);
			if (open_result < 0) {
				close_session(session);
				return unexpected_error(
					ERR_NETWORK,
					"failed to open RTMP output '" + url + "': " + av_error_string(open_result));
			}
		}

		const auto header_result = avformat_write_header(session.format_context, nullptr);
		if (header_result < 0) {
			close_session(session);
			return unexpected_error(
				ERR_PROTOCOL,
				"failed to write RTMP header for '" + url + "': " + av_error_string(header_result));
		}

		spdlog::info(
			"RTMP_OUTPUT_READY backend=libavformat codec={} url={}",
			to_string(stream_info.codec),
			url);
		return session;
	}

	static void close_session(Session &session) {
		if (session.format_context == nullptr) {
			return;
		}

		av_write_trailer(session.format_context);
		if (!(session.format_context->oformat->flags & AVFMT_NOFILE) && session.format_context->pb != nullptr) {
			avio_closep(&session.format_context->pb);
		}
		avformat_free_context(session.format_context);
		session.format_context = nullptr;
		session.video_stream   = nullptr;
	}

	CodecType codec_{CodecType::H264};
	std::vector<Session> sessions_{};
	std::uint64_t access_units_{0};
	std::uint64_t video_messages_{0};
	std::uint64_t bytes_sent_{0};
};

class FfmpegProcessRtmpOutput {
public:
	struct Session {
		std::string url{};
		pid_t pid{-1};
		int stdin_fd{-1};
	};

	FfmpegProcessRtmpOutput() = default;
	FfmpegProcessRtmpOutput(const FfmpegProcessRtmpOutput &) = delete;
	FfmpegProcessRtmpOutput &operator=(const FfmpegProcessRtmpOutput &) = delete;
	FfmpegProcessRtmpOutput(FfmpegProcessRtmpOutput &&) noexcept = default;
	FfmpegProcessRtmpOutput &operator=(FfmpegProcessRtmpOutput &&) noexcept = default;

	~FfmpegProcessRtmpOutput() {
		for (auto &session : sessions_) {
			close_session(session);
		}
	}

	[[nodiscard]]
	static Result<RtmpOutput> create(
		const RuntimeConfig &config,
		const encode::EncodedStreamInfo &stream_info) {
		FfmpegProcessRtmpOutput output{};
		output.codec_       = stream_info.codec;
		output.ffmpeg_path_ = config.outputs.rtmp.ffmpeg_path;
		output.frame_rate_  = AVRational{
			static_cast<int>(stream_info.frame_rate_num),
			static_cast<int>(stream_info.frame_rate_den),
		};
		auto decoder_config_annexb = decoder_config_to_annexb(stream_info.codec, stream_info.decoder_config);
		if (!decoder_config_annexb) {
			return unexpected_error(
				ERR_PROTOCOL,
				"ffmpeg_process RTMP requires decoder config: " + format_error(decoder_config_annexb.error()));
		}
		output.decoder_config_annexb_ = std::move(*decoder_config_annexb);

		for (const auto &url : config.outputs.rtmp.urls) {
			auto session = output.spawn_session(url);
			if (!session) {
				return std::unexpected(session.error());
			}
			output.sessions_.push_back(std::move(*session));
		}

		return pro::make_proxy<RtmpOutputFacade, FfmpegProcessRtmpOutput>(std::move(output));
	}

	[[nodiscard]]
	std::string_view backend_name() const {
		return "ffmpeg_process";
	}

	[[nodiscard]]
	Status publish_access_unit(const encode::EncodedAccessUnit &access_unit) {
		for (auto &session : sessions_) {
			auto process_state = poll_session(session);
			if (!process_state) {
				return std::unexpected(process_state.error());
			}

			auto write_result = write_all(session.stdin_fd, access_unit.annexb_bytes);
			if (!write_result) {
				return unexpected_error(
					ERR_IO,
					"ffmpeg_process RTMP write failed for '" + session.url + "': " + format_error(write_result.error()));
			}

			video_messages_ += 1;
			bytes_sent_ += access_unit.annexb_bytes.size();
		}

		access_units_ += 1;
		return {};
	}

	void log_metrics() const {
		spdlog::info(
			"RTMP_OUTPUT_METRICS backend={} codec={} urls={} access_units={} video_messages={} bytes_sent={} ffmpeg_path={}",
			backend_name(),
			to_string(codec_),
			sessions_.size(),
			access_units_,
			video_messages_,
			bytes_sent_,
			ffmpeg_path_);
	}

private:
	[[nodiscard]]
	Result<Session> spawn_session(const std::string &url) const {
		int stdin_pipe[2]{-1, -1};
		if (pipe(stdin_pipe) != 0) {
			return unexpected_error(ERR_IO, "failed to create ffmpeg stdin pipe: " + std::string(std::strerror(errno)));
		}

		const auto child = fork();
		if (child < 0) {
			close(stdin_pipe[0]);
			close(stdin_pipe[1]);
			return unexpected_error(ERR_CHILD_PROCESS, "failed to fork ffmpeg child: " + std::string(std::strerror(errno)));
		}

		if (child == 0) {
			dup2(stdin_pipe[0], STDIN_FILENO);
			close(stdin_pipe[0]);
			close(stdin_pipe[1]);

			const int null_fd = open("/dev/null", O_WRONLY);
			if (null_fd >= 0) {
				dup2(null_fd, STDOUT_FILENO);
				close(null_fd);
			}

			const std::string frame_rate =
				std::to_string(frame_rate_.num) + "/" + std::to_string(std::max(frame_rate_.den, 1));

			std::vector<std::string> args_storage{
				ffmpeg_path_,
				"-hide_banner",
				"-loglevel",
				"warning",
				"-fflags",
				"+genpts+nobuffer",
				"-use_wallclock_as_timestamps",
				"1",
				"-framerate",
				frame_rate,
				"-f",
				ffmpeg_input_format(codec_),
				"-i",
				"pipe:0",
				"-an",
				"-c:v",
				"copy",
				"-f",
				"flv",
				url,
			};

			std::vector<char *> argv{};
			argv.reserve(args_storage.size() + 1);
			for (auto &arg : args_storage) {
				argv.push_back(arg.data());
			}
			argv.push_back(nullptr);

			execvp(argv[0], argv.data());
			::_exit(127);
		}

		close(stdin_pipe[0]);

		Session session{};
		session.url      = url;
		session.pid      = child;
		session.stdin_fd = stdin_pipe[1];

		auto preamble_write = write_all(
			session.stdin_fd,
			std::span<const std::uint8_t>(decoder_config_annexb_.data(), decoder_config_annexb_.size()));
		if (!preamble_write) {
			close_session(session);
			return unexpected_error(
				ERR_IO,
				"failed to seed ffmpeg_process decoder config: " + format_error(preamble_write.error()));
		}

		spdlog::info(
			"RTMP_OUTPUT_READY backend=ffmpeg_process codec={} url={} ffmpeg_path={} pid={}",
			to_string(codec_),
			url,
			ffmpeg_path_,
			static_cast<long long>(session.pid));
		return session;
	}

	[[nodiscard]]
	static Status poll_session(Session &session) {
		if (session.pid <= 0) {
			return unexpected_error(ERR_CHILD_PROCESS, "ffmpeg child is not running");
		}

		int status{0};
		const auto wait_result = waitpid(session.pid, &status, WNOHANG);
		if (wait_result == 0) {
			return {};
		}
		if (wait_result < 0) {
			return unexpected_error(ERR_CHILD_PROCESS, "failed to poll ffmpeg child: " + std::string(std::strerror(errno)));
		}

		if (session.stdin_fd >= 0) {
			close(session.stdin_fd);
			session.stdin_fd = -1;
		}
		session.pid = -1;
		return unexpected_error(ERR_CHILD_PROCESS, "ffmpeg child exited before publish completed");
	}

	static void close_session(Session &session) {
		if (session.stdin_fd >= 0) {
			close(session.stdin_fd);
			session.stdin_fd = -1;
		}
		if (session.pid > 0) {
			kill(session.pid, SIGTERM);
			(void)waitpid(session.pid, nullptr, 0);
			session.pid = -1;
		}
	}

	CodecType codec_{CodecType::H264};
	std::string ffmpeg_path_{};
	AVRational frame_rate_{30, 1};
	std::vector<std::uint8_t> decoder_config_annexb_{};
	std::vector<Session> sessions_{};
	std::uint64_t access_units_{0};
	std::uint64_t video_messages_{0};
	std::uint64_t bytes_sent_{0};
};

}

Result<RtmpOutput> make_rtmp_output(
	const RuntimeConfig &config,
	const encode::EncodedStreamInfo &stream_info) {
	switch (config.outputs.rtmp.transport) {
	case RtmpTransportType::Libavformat:
		return LibavformatRtmpOutput::create(config, stream_info);
	case RtmpTransportType::FfmpegProcess:
		return FfmpegProcessRtmpOutput::create(config, stream_info);
	case RtmpTransportType::LegacyCustom: {
		auto publisher = RtmpPublisher::create(config);
		if (!publisher) {
			return unexpected_error(ERR_PROTOCOL, publisher.error());
		}
		return pro::make_proxy<RtmpOutputFacade, LegacyCustomRtmpOutput>(std::move(*publisher));
	}
	}

	return unexpected_error(ERR_INTERNAL, "unknown RTMP transport");
}

}