#include <array>
#include <cerrno>
#include <chrono>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <expected>
#include <fstream>
#include <optional>
#include <sstream>
#include <string>
#include <string_view>
#include <thread>
#include <vector>

#include <netinet/in.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <unistd.h>

#include <spdlog/spdlog.h>

#include "cvmmap_streamer/common.h"

namespace {

// RFC3550 RTP header constants
constexpr std::size_t kRtpHeaderMinSize    = 12;
constexpr std::uint8_t kRtpVersion         = 2;
constexpr std::uint8_t kRtpVersionMask     = 0xC0;
constexpr std::uint8_t kRtpVersionShift    = 6;
constexpr std::uint8_t kRtpPaddingMask     = 0x20;
constexpr std::uint8_t kRtpExtensionMask   = 0x10;
constexpr std::uint8_t kRtpCsrcCountMask   = 0x0F;
constexpr std::uint8_t kRtpMarkerMask      = 0x80;
constexpr std::uint8_t kRtpPayloadTypeMask = 0x7F;

// RTP header structure (RFC3550)
struct RtpHeader {
	std::uint8_t version;     // 2 bits
	bool padding;             // 1 bit
	bool extension;           // 1 bit
	std::uint8_t csrcCount;   // 4 bits
	bool marker;              // 1 bit
	std::uint8_t payloadType; // 7 bits
	std::uint16_t sequence;   // 16 bits
	std::uint32_t timestamp;  // 32 bits
	std::uint32_t ssrc;       // 32 bits
};

// Parsed SDP media info
struct SdpMediaInfo {
	std::string encodingName;
	std::uint32_t clockRate  = 0;
	std::uint8_t payloadType = 0;
	bool hasH264             = false;
	bool hasH265             = false;
};

// Test configuration
struct Config {
	std::uint16_t port = 5004;
	std::optional<std::uint8_t> expectedPt;
	std::optional<std::string> sdpFile;
	std::optional<std::string> decodeHook;
	std::uint32_t packetThreshold = 10;
	std::uint32_t timeoutMs       = 5000;
	bool verbose                  = false;
};

// Test statistics
struct Stats {
	std::uint32_t packetsReceived = 0;
	std::uint32_t sequenceGaps    = 0;
	std::uint32_t invalidPackets  = 0;
	std::uint16_t lastSequence    = 0;
	std::uint8_t detectedPt       = 0;
	bool hasSeenPacket            = false;
	std::optional<std::uint8_t> ptMismatchError;
};

// Parse RTP header from buffer
std::optional<RtpHeader> parseRtpHeader(std::span<const std::uint8_t> data) {
	if (data.size() < kRtpHeaderMinSize) {
		return std::nullopt;
	}

	RtpHeader header{};
	header.version     = (data[0] & kRtpVersionMask) >> kRtpVersionShift;
	header.padding     = (data[0] & kRtpPaddingMask) != 0;
	header.extension   = (data[0] & kRtpExtensionMask) != 0;
	header.csrcCount   = data[0] & kRtpCsrcCountMask;
	header.marker      = (data[1] & kRtpMarkerMask) != 0;
	header.payloadType = data[1] & kRtpPayloadTypeMask;
	header.sequence    = static_cast<std::uint16_t>(data[2]) << 8 | data[3];
	header.timestamp   = static_cast<std::uint32_t>(data[4]) << 24 |
					   static_cast<std::uint32_t>(data[5]) << 16 |
					   static_cast<std::uint32_t>(data[6]) << 8 | data[7];
	header.ssrc = static_cast<std::uint32_t>(data[8]) << 24 |
				  static_cast<std::uint32_t>(data[9]) << 16 |
				  static_cast<std::uint32_t>(data[10]) << 8 | data[11];

	return header;
}

// Parse SDP file for media information
std::optional<SdpMediaInfo> parseSdpFile(std::string_view path) {
	std::string pathStr(path);
	std::ifstream file(pathStr);
	if (!file.is_open()) {
		spdlog::error("Failed to open SDP file: {}", path);
		return std::nullopt;
	}

	SdpMediaInfo info;
	std::string line;
	bool inMediaSection = false;

	while (std::getline(file, line)) {
		// Remove trailing \r if present
		if (!line.empty() && line.back() == '\r') {
			line.pop_back();
		}

		if (line.starts_with("m=")) {
			inMediaSection = true;
			// Parse media line: m=<media> <port> <proto> <pt>
			// e.g., m=video 5004 RTP/AVP 96
			std::istringstream iss(line);
			std::string mediaType, port, proto;
			int pt = 0;
			// Skip "m=" prefix
			if (line.size() > 2 && iss.seekg(2) && std::getline(iss, mediaType, ' ')) {
				if (iss >> port >> proto >> pt) {
					info.payloadType = static_cast<std::uint8_t>(pt);
				}
			}
		} else if (inMediaSection && line.starts_with("a=rtpmap:")) {
			// Parse rtpmap: a=rtpmap:<pt> <encoding>/<clockrate>
			// e.g., a=rtpmap:96 H264/90000
			size_t ptEnd = line.find(' ', 9);
			if (ptEnd != std::string::npos) {
				size_t slashPos = line.find('/', ptEnd + 1);
				if (slashPos != std::string::npos) {
					info.encodingName = line.substr(ptEnd + 1, slashPos - ptEnd - 1);
					info.clockRate    = std::stoul(line.substr(slashPos + 1));
				}
			}
			if (info.encodingName == "H264") {
				info.hasH264 = true;
			} else if (info.encodingName == "H265" || info.encodingName == "HEVC") {
				info.hasH265 = true;
			}
		}
	}

	return info;
}

// Create UDP socket and bind to port
std::expected<int, std::string> createUdpSocket(std::uint16_t port) {
	int sock = socket(AF_INET, SOCK_DGRAM, 0);
	if (sock < 0) {
		return std::unexpected(std::format("socket failed: {}", std::strerror(errno)));
	}

	int reuse = 1;
	if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0) {
		close(sock);
		return std::unexpected(std::format("setsockopt failed: {}", std::strerror(errno)));
	}

	sockaddr_in addr{};
	addr.sin_family      = AF_INET;
	addr.sin_port        = htons(port);
	addr.sin_addr.s_addr = INADDR_ANY;

	if (bind(sock, reinterpret_cast<sockaddr *>(&addr), sizeof(addr)) < 0) {
		close(sock);
		return std::unexpected(std::format("bind failed: {}", std::strerror(errno)));
	}

	return sock;
}

// Parse command-line arguments
std::expected<Config, std::string> parseArgs(int argc, char **argv) {
	Config config;

	for (int i = 1; i < argc; ++i) {
		std::string_view arg(argv[i]);

		if (arg == "--help" || arg == "-h") {
			return std::unexpected("help");
		} else if (arg == "--port" && i + 1 < argc) {
			config.port = static_cast<std::uint16_t>(std::stoul(argv[++i]));
		} else if (arg == "--expect-pt" && i + 1 < argc) {
			config.expectedPt = static_cast<std::uint8_t>(std::stoul(argv[++i]));
		} else if (arg == "--sdp" && i + 1 < argc) {
			config.sdpFile = argv[++i];
		} else if (arg == "--decode-hook" && i + 1 < argc) {
			config.decodeHook = argv[++i];
		} else if (arg == "--packet-threshold" && i + 1 < argc) {
			config.packetThreshold = std::stoul(argv[++i]);
		} else if (arg == "--timeout-ms" && i + 1 < argc) {
			config.timeoutMs = std::stoul(argv[++i]);
		} else if (arg == "--verbose" || arg == "-v") {
			config.verbose = true;
		}
	}

	return config;
}

// Print usage
void printRtpReceiverUsage() {
	spdlog::info("rtp_receiver_tester - UDP RTP packet receiver and validator");
	spdlog::info("");
	spdlog::info("Usage:");
	spdlog::info("  rtp_receiver_tester [options]");
	spdlog::info("");
	spdlog::info("Options:");
	spdlog::info("  --port <num>           UDP port to listen on (default: 5004)");
	spdlog::info("  --expect-pt <num>      Expected payload type (0-127)");
	spdlog::info("  --sdp <path>           SDP file to validate against");
	spdlog::info("  --decode-hook <cmd>    Optional command to validate payload");
	spdlog::info("  --packet-threshold <n> Minimum packets to consider success (default: 10)");
	spdlog::info("  --timeout-ms <ms>      Max time to wait for packets (default: 5000)");
	spdlog::info("  --verbose, -v          Enable verbose logging");
	spdlog::info("  --help, -h             Show this message");
	spdlog::info("");
	spdlog::info("Examples:");
	spdlog::info("  rtp_receiver_tester --port 5004 --expect-pt 96");
	spdlog::info("  rtp_receiver_tester --port 5004 --sdp /tmp/stream.sdp");
	spdlog::info("");
	spdlog::info("Exit codes:");
	spdlog::info("  0  Success (packets received, PT matches)");
	spdlog::info("  1  Invalid arguments");
	spdlog::info("  2  Socket/bind error");
	spdlog::info("  3  Payload type mismatch");
	spdlog::info("  4  Packet threshold not met");
	spdlog::info("  5  SDP validation failed");
	spdlog::info("  6  Decode hook failed");
}

// Run optional decode hook
bool runDecodeHook(std::string_view hookCmd, std::span<const std::uint8_t> payload) {
	if (hookCmd.empty()) {
		return true;
	}

	pid_t pid = fork();
	if (pid < 0) {
		spdlog::error("fork failed for decode hook");
		return false;
	}

	if (pid == 0) {
		// Child process - execute hook
		// Write payload to stdin of hook command
		// For simplicity, use a temp file approach or pipe
		// Here we use execlp with the command
		execlp("sh", "sh", "-c", std::string(hookCmd).c_str(), nullptr);
		_exit(127);
	}

	// Parent - wait for child with timeout
	int status     = 0;
	int waitResult = waitpid(pid, &status, WNOHANG);

	// Simple non-blocking check - if not ready, we continue
	// The hook is optional/validation only
	if (waitResult == 0) {
		// Still running, don't block
		spdlog::debug("Decode hook still running (non-blocking)");
		return true;
	}

	if (WIFEXITED(status) && WEXITSTATUS(status) == 0) {
		return true;
	}

	spdlog::warn("Decode hook exited with error");
	return false;
}

} // namespace

int main(int argc, char **argv) {
	if (argc <= 1 || cvmmap_streamer::has_help_flag(argc, argv)) {
		printRtpReceiverUsage();
		return (argc <= 1) ? 1 : 0;
	}

	auto configResult = parseArgs(argc, argv);
	if (!configResult) {
		if (configResult.error() == "help") {
			printRtpReceiverUsage();
			return 0;
		}
		spdlog::error("Argument error: {}", configResult.error());
		printRtpReceiverUsage();
		return 1;
	}

	const auto &config = *configResult;

	if (config.verbose) {
		spdlog::set_level(spdlog::level::debug);
	}

	// Parse SDP file if provided
	std::optional<SdpMediaInfo> sdpInfo;
	if (config.sdpFile) {
		sdpInfo = parseSdpFile(*config.sdpFile);
		if (!sdpInfo) {
			spdlog::error("Failed to parse SDP file: {}", *config.sdpFile);
			return 5;
		}
		spdlog::info("SDP parsed: encoding={}, clock-rate={}, PT={}",
					 sdpInfo->encodingName,
					 sdpInfo->clockRate,
					 sdpInfo->payloadType);

		// Cross-validate expected PT with SDP
		if (config.expectedPt && *config.expectedPt != sdpInfo->payloadType) {
			spdlog::error("Expected PT({}) does not match SDP PT({})",
						  *config.expectedPt,
						  sdpInfo->payloadType);
			return 5;
		}
	}

	// Create UDP socket
	auto sockResult = createUdpSocket(config.port);
	if (!sockResult) {
		spdlog::error("Socket error: {}", sockResult.error());
		return 2;
	}

	int sock = *sockResult;
	spdlog::info("Listening on UDP port {} for RTP packets...", config.port);

	Stats stats;
	auto startTime = std::chrono::steady_clock::now();

	std::vector<std::uint8_t> buffer(65535);
	sockaddr_in clientAddr{};
	socklen_t addrLen = sizeof(clientAddr);

	while (true) {
		// Check timeout
		auto elapsed = std::chrono::steady_clock::now() - startTime;
		if (std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count() > config.timeoutMs) {
			spdlog::info("Timeout reached after {} ms", config.timeoutMs);
			break;
		}

		// Non-blocking receive with short timeout using select
		fd_set readfds;
		FD_ZERO(&readfds);
		FD_SET(sock, &readfds);

		timeval tv{};
		tv.tv_sec  = 0;
		tv.tv_usec = 100000; // 100ms

		int selectResult = select(sock + 1, &readfds, nullptr, nullptr, &tv);
		if (selectResult < 0) {
			spdlog::error("select error: {}", std::strerror(errno));
			break;
		}
		if (selectResult == 0) {
			continue; // Timeout, check overall timeout
		}

		ssize_t received = recvfrom(sock,
									buffer.data(),
									buffer.size(),
									0,
									reinterpret_cast<sockaddr *>(&clientAddr),
									&addrLen);

		if (received < 0) {
			if (errno == EAGAIN || errno == EWOULDBLOCK) {
				continue;
			}
			spdlog::error("recvfrom error: {}", std::strerror(errno));
			break;
		}

		// Parse RTP header
		auto headerOpt = parseRtpHeader(std::span(buffer.data(), received));
		if (!headerOpt) {
			spdlog::warn("Received invalid packet (too small)");
			stats.invalidPackets++;
			continue;
		}

		const auto &header = *headerOpt;

		// Validate RTP version
		if (header.version != kRtpVersion) {
			spdlog::warn("Invalid RTP version: {} (expected {})", header.version, kRtpVersion);
			stats.invalidPackets++;
			continue;
		}

		stats.packetsReceived++;

		// Track sequence gaps
		if (stats.hasSeenPacket) {
			std::uint16_t expectedSeq = stats.lastSequence + 1;
			if (header.sequence != expectedSeq) {
				std::uint16_t gap = header.sequence - expectedSeq;
				stats.sequenceGaps += gap;
				if (config.verbose) {
					spdlog::debug("Sequence gap detected: expected {}, got {} (gap size: {})",
								  expectedSeq,
								  header.sequence,
								  gap);
				}
			}
		}
		stats.lastSequence  = header.sequence;
		stats.hasSeenPacket = true;
		stats.detectedPt    = header.payloadType;

		// Validate payload type if expected
		if (config.expectedPt && header.payloadType != *config.expectedPt) {
			spdlog::error("Payload type mismatch: expected {}, got {}",
						  *config.expectedPt,
						  header.payloadType);
			stats.ptMismatchError = header.payloadType;
		}

		// Log packet info
		if (config.verbose) {
			spdlog::debug("Packet {}: PT={}, seq={}, ts={}, ssrc=0x{:08X}, marker={}",
						  stats.packetsReceived,
						  header.payloadType,
						  header.sequence,
						  header.timestamp,
						  header.ssrc,
						  header.marker);
		}

		// Run optional decode hook
		if (config.decodeHook && !stats.ptMismatchError) {
			std::size_t headerSize = kRtpHeaderMinSize + (header.csrcCount * 4);
			if (header.extension) {
				// Skip extension header if present
				if (received >= headerSize + 4) {
					std::uint16_t extLength = static_cast<std::uint16_t>(buffer[headerSize + 2]) << 8 |
											  buffer[headerSize + 3];
					headerSize += 4 + (extLength * 4);
				}
			}
			if (received > static_cast<ssize_t>(headerSize)) {
				auto payload = std::span(buffer.data() + headerSize, received - headerSize);
				if (!runDecodeHook(*config.decodeHook, payload)) {
					spdlog::warn("Decode hook validation failed");
				}
			}
		}

		// Check if we've received enough packets
		if (stats.packetsReceived >= config.packetThreshold) {
			spdlog::info("Packet threshold reached ({} packets)", config.packetThreshold);
			break;
		}
	}

	close(sock);

	// Print summary
	spdlog::info("");
	spdlog::info("=== RTP Receiver Statistics ===");
	spdlog::info("Packets received: {}", stats.packetsReceived);
	spdlog::info("Sequence gaps: {}", stats.sequenceGaps);
	spdlog::info("Invalid packets: {}", stats.invalidPackets);
	if (stats.hasSeenPacket) {
		spdlog::info("Detected payload type: {}", stats.detectedPt);
	}

	// Determine exit code
	if (stats.ptMismatchError) {
		spdlog::error("FAIL: Payload type mismatch detected (expected {}, got {})",
					  config.expectedPt.value(),
					  *stats.ptMismatchError);
		return 3;
	}

	if (stats.packetsReceived < config.packetThreshold) {
		spdlog::error("FAIL: Packet threshold not met (received {}, required {})",
					  stats.packetsReceived,
					  config.packetThreshold);
		return 4;
	}

	spdlog::info("PASS: All validations successful");
	return 0;
}