#include <mcap/reader.hpp>

#include "cvmmap_streamer/DepthMap.pb.h"
#include "cvmmap_streamer/common.h"
#include "cvmmap_streamer/record/mcap_record_sink.hpp"
#include "foxglove/CompressedVideo.pb.h"

#include <rvl/rvl.hpp>
#include <spdlog/spdlog.h>

#include <cmath>
#include <cstdint>
#include <filesystem>
#include <limits>
#include <optional>
#include <string>
#include <vector>

namespace {

enum class TesterExitCode : int {
	Success = 0,
	CreateError = 2,
	WriteError = 3,
	OpenError = 4,
	VerificationError = 5,
};

[[nodiscard]]
constexpr int exit_code(TesterExitCode code) {
	return static_cast<int>(code);
}

[[nodiscard]]
bool approx_equal(float lhs, float rhs, float tolerance) {
	return std::fabs(lhs - rhs) <= tolerance;
}

}

int main(int argc, char **argv) {
	if (cvmmap_streamer::has_help_flag(argc, argv)) {
		cvmmap_streamer::print_help("mcap_depth_record_tester");
		return exit_code(TesterExitCode::Success);
	}

	const std::filesystem::path output_path =
		argc > 1
			? std::filesystem::path(argv[1])
			: std::filesystem::temp_directory_path() / "cvmmap_streamer_depth_record_test.mcap";
	if (output_path.has_parent_path()) {
		std::filesystem::create_directories(output_path.parent_path());
	}

	cvmmap_streamer::RuntimeConfig config = cvmmap_streamer::RuntimeConfig::defaults();
	config.record.mcap.enabled = true;
	config.record.mcap.path = output_path.string();
	config.record.mcap.topic = "/camera/video";
	config.record.mcap.depth_topic = "/camera/depth";
	config.record.mcap.frame_id = "camera";
	config.record.mcap.compression = cvmmap_streamer::McapCompression::None;

	cvmmap_streamer::encode::EncodedStreamInfo stream_info{};
	stream_info.codec = cvmmap_streamer::CodecType::H264;

	auto sink = cvmmap_streamer::record::McapRecordSink::create(config, stream_info);
	if (!sink) {
		spdlog::error("failed to create MCAP sink: {}", sink.error());
		return exit_code(TesterExitCode::CreateError);
	}

	cvmmap_streamer::encode::EncodedAccessUnit access_unit{};
	access_unit.codec = cvmmap_streamer::CodecType::H264;
	access_unit.source_timestamp_ns = 10;
	access_unit.stream_pts_ns = 10;
	access_unit.keyframe = false;
	access_unit.annexb_bytes = {0x00, 0x00, 0x00, 0x01, 0x09, 0x10};
	if (auto write = sink->write_access_unit(access_unit); !write) {
		spdlog::error("failed to write video access unit: {}", write.error());
		return exit_code(TesterExitCode::WriteError);
	}

	const float depth_mm_pixels[4] = {
		1000.0f,
		2000.0f,
		std::numeric_limits<float>::quiet_NaN(),
		0.0f,
	};
	cvmmap_streamer::record::RawDepthMapView depth_mm{};
	depth_mm.timestamp_ns = 11;
	depth_mm.width = 2;
	depth_mm.height = 2;
	depth_mm.source_unit = cvmmap_streamer::ipc::DepthUnit::Millimeter;
	depth_mm.pixels = depth_mm_pixels;
	if (auto write = sink->write_depth_map(depth_mm); !write) {
		spdlog::error("failed to write millimeter depth map: {}", write.error());
		return exit_code(TesterExitCode::WriteError);
	}

	const float depth_m_pixels[4] = {
		1.0f,
		2.0f,
		std::numeric_limits<float>::quiet_NaN(),
		0.0f,
	};
	cvmmap_streamer::record::RawDepthMapView depth_m{};
	depth_m.timestamp_ns = 12;
	depth_m.width = 2;
	depth_m.height = 2;
	depth_m.source_unit = cvmmap_streamer::ipc::DepthUnit::Meter;
	depth_m.pixels = depth_m_pixels;
	if (auto write = sink->write_depth_map(depth_m); !write) {
		spdlog::error("failed to write meter depth map: {}", write.error());
		return exit_code(TesterExitCode::WriteError);
	}

	const float depth_unknown_pixels[4] = {
		1500.0f,
		2500.0f,
		std::numeric_limits<float>::quiet_NaN(),
		0.0f,
	};
	cvmmap_streamer::record::RawDepthMapView depth_unknown{};
	depth_unknown.timestamp_ns = 13;
	depth_unknown.width = 2;
	depth_unknown.height = 2;
	depth_unknown.source_unit = cvmmap_streamer::ipc::DepthUnit::Unknown;
	depth_unknown.pixels = depth_unknown_pixels;
	if (auto write = sink->write_depth_map(depth_unknown); !write) {
		spdlog::error("failed to write unknown-unit depth map: {}", write.error());
		return exit_code(TesterExitCode::WriteError);
	}

	sink->close();

	mcap::McapReader reader{};
	const auto open_status = reader.open(output_path.string());
	if (!open_status.ok()) {
		spdlog::error("failed to open MCAP file '{}': {}", output_path.string(), open_status.message);
		return exit_code(TesterExitCode::OpenError);
	}

	std::uint64_t video_messages{0};
	std::vector<cvmmap_streamer::DepthMap> depth_messages{};

	auto messages = reader.readMessages();
	for (auto it = messages.begin(); it != messages.end(); ++it) {
		if (it->schema == nullptr || it->channel == nullptr) {
			spdlog::error("MCAP message missing schema or channel");
			reader.close();
			return exit_code(TesterExitCode::VerificationError);
		}
		if (it->schema->encoding != "protobuf" || it->channel->messageEncoding != "protobuf") {
			spdlog::error("unexpected schema encoding in MCAP file");
			reader.close();
			return exit_code(TesterExitCode::VerificationError);
		}

		if (it->schema->name == "foxglove.CompressedVideo") {
			foxglove::CompressedVideo video{};
			if (!video.ParseFromArray(it->message.data, static_cast<int>(it->message.dataSize))) {
				spdlog::error("failed to parse foxglove.CompressedVideo payload");
				reader.close();
				return exit_code(TesterExitCode::VerificationError);
			}
			if (it->channel->topic != "/camera/video" || video.frame_id() != "camera" || video.data().empty()) {
				spdlog::error("video MCAP payload verification failed");
				reader.close();
				return exit_code(TesterExitCode::VerificationError);
			}
			video_messages += 1;
			continue;
		}

		if (it->schema->name == "cvmmap_streamer.DepthMap") {
			cvmmap_streamer::DepthMap depth{};
			if (!depth.ParseFromArray(it->message.data, static_cast<int>(it->message.dataSize))) {
				spdlog::error("failed to parse cvmmap_streamer.DepthMap payload");
				reader.close();
				return exit_code(TesterExitCode::VerificationError);
			}
			if (it->channel->topic != "/camera/depth" || depth.frame_id() != "camera") {
				spdlog::error("depth MCAP payload verification failed");
				reader.close();
				return exit_code(TesterExitCode::VerificationError);
			}
			depth_messages.push_back(std::move(depth));
		}
	}

	reader.close();

	if (video_messages != 1 || depth_messages.size() != 3) {
		spdlog::error(
			"unexpected message counts: video={} depth={}",
			video_messages,
			depth_messages.size());
		return exit_code(TesterExitCode::VerificationError);
	}

	const auto &mm_message = depth_messages[0];
	if (mm_message.source_unit() != cvmmap_streamer::DepthMap::DEPTH_UNIT_MILLIMETER ||
		mm_message.storage_unit() != cvmmap_streamer::DepthMap::STORAGE_UNIT_MILLIMETER ||
		mm_message.encoding() != cvmmap_streamer::DepthMap::RVL_U16_LOSSLESS) {
		spdlog::error("millimeter depth metadata verification failed");
		return exit_code(TesterExitCode::VerificationError);
	}
	const auto mm_info = rvl::inspect_image(std::span<const std::uint8_t>(
		reinterpret_cast<const std::uint8_t *>(mm_message.data().data()),
		mm_message.data().size()));
	const auto mm_decoded = rvl::decompress_image(std::span<const std::uint8_t>(
		reinterpret_cast<const std::uint8_t *>(mm_message.data().data()),
		mm_message.data().size()));
	if (mm_info.format != rvl::ImageFormat::UInt16Lossless ||
		mm_info.rows != 2 ||
		mm_info.cols != 2 ||
		mm_decoded.pixels.size() != 4 ||
		mm_decoded.pixels[0] != 1000 ||
		mm_decoded.pixels[1] != 2000 ||
		mm_decoded.pixels[2] != 0 ||
		mm_decoded.pixels[3] != 0) {
		spdlog::error("millimeter RVL round-trip verification failed");
		return exit_code(TesterExitCode::VerificationError);
	}

	const auto &m_message = depth_messages[1];
	if (m_message.source_unit() != cvmmap_streamer::DepthMap::DEPTH_UNIT_METER ||
		m_message.storage_unit() != cvmmap_streamer::DepthMap::STORAGE_UNIT_METER ||
		m_message.encoding() != cvmmap_streamer::DepthMap::RVL_F32) {
		spdlog::error("meter depth metadata verification failed");
		return exit_code(TesterExitCode::VerificationError);
	}
	const auto m_info = rvl::inspect_image(std::span<const std::uint8_t>(
		reinterpret_cast<const std::uint8_t *>(m_message.data().data()),
		m_message.data().size()));
	const auto m_decoded = rvl::decompress_float_image(std::span<const std::uint8_t>(
		reinterpret_cast<const std::uint8_t *>(m_message.data().data()),
		m_message.data().size()));
	if (m_info.format != rvl::ImageFormat::Float32InverseDepth ||
		m_info.rows != 2 ||
		m_info.cols != 2 ||
		m_decoded.pixels.size() != 4 ||
		!approx_equal(m_decoded.pixels[0], 1.0f, 0.02f) ||
		!approx_equal(m_decoded.pixels[1], 2.0f, 0.02f) ||
		!std::isnan(m_decoded.pixels[2]) ||
		!std::isnan(m_decoded.pixels[3])) {
		spdlog::error("meter RVL round-trip verification failed");
		return exit_code(TesterExitCode::VerificationError);
	}

	const auto &unknown_message = depth_messages[2];
	if (unknown_message.source_unit() != cvmmap_streamer::DepthMap::DEPTH_UNIT_MILLIMETER ||
		unknown_message.storage_unit() != cvmmap_streamer::DepthMap::STORAGE_UNIT_MILLIMETER ||
		unknown_message.encoding() != cvmmap_streamer::DepthMap::RVL_U16_LOSSLESS) {
		spdlog::error("unknown-unit fallback metadata verification failed");
		return exit_code(TesterExitCode::VerificationError);
	}
	const auto unknown_info = rvl::inspect_image(std::span<const std::uint8_t>(
		reinterpret_cast<const std::uint8_t *>(unknown_message.data().data()),
		unknown_message.data().size()));
	const auto unknown_decoded = rvl::decompress_image(std::span<const std::uint8_t>(
		reinterpret_cast<const std::uint8_t *>(unknown_message.data().data()),
		unknown_message.data().size()));
	if (unknown_info.format != rvl::ImageFormat::UInt16Lossless ||
		unknown_info.rows != 2 ||
		unknown_info.cols != 2 ||
		unknown_decoded.pixels.size() != 4 ||
		unknown_decoded.pixels[0] != 1500 ||
		unknown_decoded.pixels[1] != 2500 ||
		unknown_decoded.pixels[2] != 0 ||
		unknown_decoded.pixels[3] != 0) {
		spdlog::error("unknown-unit fallback RVL round-trip verification failed");
		return exit_code(TesterExitCode::VerificationError);
	}

	spdlog::info(
		"validated same-file MCAP video+depth recording at '{}'",
		output_path.string());
	return exit_code(TesterExitCode::Success);
}