// clang-format off
#pragma comment(lib, "ws2_32.lib")
#include <WinSock2.h>
#include <ws2tcpip.h>

#include <windows.h>
#include <stdio.h>
#include <string>
#include <vector>

#include "capture.hpp"
#include "Utils.hpp"
// clang-format on

#pragma pack(push, 1)
enum class MessageType : u32 {
	Resize, // tResizeMessage
	Data,	// tDataMessage
};

struct tMessageHeader {
	MessageType type;
	u32 datalen;
	// data
};

struct tResizeMessage {
	u32 width;
	u32 height;
};

struct tTile {
	u32 tile_x;
	u32 tile_y;

	u32 tile_width;
	u32 tile_height;

	// u32 tile_rgba[tile_width*tile_height]
};

struct tDataMessage {
	u32 tileCount;
};

#pragma pack(pop)

struct tileRect {
	u32 x, y, width, height;
};

// client for streamserver
class cStreamClient {
	SOCKET tcpSocket { -1 };

   public:
	cStreamClient() = default;

	~cStreamClient() {
		if(tcpSocket != -1) {
			closesocket(tcpSocket);
			tcpSocket = -1;
		}
	}

	bool Connect(const char* address, int port) {
		tcpSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
		if(tcpSocket == -1) {
			return false;
		}

		sockaddr_in clientSvc {};
		clientSvc.sin_family = AF_INET;
		inet_pton(AF_INET, address, &clientSvc.sin_addr.s_addr);
		clientSvc.sin_port = htons(port);

		if(connect(tcpSocket, (SOCKADDR*)&clientSvc, sizeof(clientSvc)) == SOCKET_ERROR) {
			printf("No connection socket. Fuck you\n");
			return false;
		}

		// disable Nagle's alrogithm
		int nodelay = 1;
		setsockopt(tcpSocket, IPPROTO_TCP, TCP_NODELAY, (const char*)&nodelay, sizeof(nodelay));
		return true;
	}

	void SendResize(const tResizeMessage& resize) {
		tMessageHeader header;
		header.type = MessageType::Resize;
		header.datalen = sizeof(tResizeMessage);

		send(tcpSocket, (const char*)&header, sizeof(header), 0);
		send(tcpSocket, (const char*)&resize, sizeof(resize), 0);
	}

	void SendData(const u32* pData, u32 width, u32 height, const std::vector<tileRect>& tiles) {
		tMessageHeader header;
		header.type = MessageType::Data;
		// header.datalen = data.get_size() * sizeof(UINT32);
		header.datalen = sizeof(tDataMessage);


		// send tile header
		bool sendFull = false;
		tDataMessage dm;

		if(tiles.empty()) {
			// Send the full screen as a "tile"
			dm.tileCount = 1;
			sendFull = true;
		} else {
			dm.tileCount = (u32)tiles.size();
		}

		// we have writev() at home
		// writev() at home:
		send(tcpSocket, (const char*)&header, sizeof(header), 0);
		send(tcpSocket, (const char*)&dm, sizeof(dm), 0);

		// send each tile
		if(!sendFull) {
			for(auto& tile : tiles) {
				tTile tile_wire { tile.x, tile.y, tile.width, tile.height };

				std::vector<u32> tileData;
				tileData.resize(tile.width * tile.height);

				for(u32 y = 0; y < tile.height; ++y) {
					auto* pTileLineStart = &pData[(tile.y + y) * width + tile.x];
					memcpy(&tileData[y * tile.width], pTileLineStart, tile.width * sizeof(u32));
				}

				// send header
				send(tcpSocket, (const char*)&tile_wire, (i32)sizeof(tile_wire), 0);
				send(tcpSocket, (const char*)&tileData[0], (i32)tileData.size() * 4, 0);

				// send data now
				/*for (auto y = tile.y; y < tile.y + tile.height; ++y) {
					auto* pTileLineStart = &pData[y * width + tile.x];
					send(tcpSocket, (const char*)pTileLineStart, tile.width * sizeof(UINT32), 0);
				}*/
			}
		} else {
			tTile tDummyTile { 0, 0, width, height };
			send(tcpSocket, (const char*)&tDummyTile, sizeof(tDummyTile), 0);
			send(tcpSocket, (const char*)&pData[0], (i32)((width * height) * sizeof(u32)), 0);
		}

		// send(tcpSocket, (const char*)&data.data()[0], data.get_size() * sizeof(UINT32), 0);
	}
};


int main(int argc, char** argv) {
	WSADATA data;
	if(WSAStartup(MAKEWORD(2, 2), &data) != NO_ERROR) {
		return 1;
	}

	cStreamClient client;

	if(!client.Connect("192.168.1.149", 9438)) {
		printf("conn failed\n");
		return 1;
	}

	// Create a capture interface
	auto capture = hazelnut::CreateFramebufferCapture(hazelnut::GuessBestCaptureInterface());
	if (!capture) {
		printf("Failed to create a capture interface\n");
	}

	printf("Successfully created a framebuffer capture interface\n");

	if(capture->Initialize()) {
		// Sleep 
		Sleep(100);

		bool firstFrame = true;
		std::vector<tileRect> tiles {};
		hazelnut::FramebufferInformation framebuffer{};
		hazelnut::DiffInformation diff{};

		while(true) {
			auto result = capture->CaptureFrame();
			if (result == hazelnut::DisplayCaptureResult::Ok || result == hazelnut::DisplayCaptureResult::OkButResized) {
				// Check for resize.
				if (result == hazelnut::DisplayCaptureResult::OkButResized) {
					framebuffer = capture->GetFramebufferInformation();
					diff = capture->GetDiffInformation();
					firstFrame = true;

					client.SendResize({ framebuffer.width, framebuffer.height });
				}

				tiles.clear();

				if(firstFrame == false) {
					for(u32 y = 0; y < diff.diffMapHeight; ++y) {
						for(u32 x = 0; x < diff.diffmapWidth; ++x) {
							auto& bl = diff.pDiffMap[y * diff.diffmapWidth + x];
							if(bl != 0) {
								tiles.push_back(tileRect {
								x * (framebuffer.width / diff.diffmapWidth),	// x
								y * (framebuffer.height / diff.diffMapHeight), // y
								framebuffer.width / diff.diffmapWidth,			// width
								framebuffer.height / diff.diffMapHeight		// height
								});
							}
						}
					}
				}

				// send that to the server
				client.SendData(framebuffer.pFramebuffer, framebuffer.width, framebuffer.height, tiles);

				if(firstFrame)
					firstFrame = false;
			} else {
				printf("Failed to capture\n");
				break;
			}
		}
	}

	return 0;
}