#include "SkyPortalIO.h"
#include "Misc/AES.h"

#include "HAL/UnrealMemory.h"
#include "openssl/aes.h"



DEFINE_LOG_CATEGORY(LogHIDApi);
DEFINE_LOG_CATEGORY(LogSkyportalIO);


//Constructor
USkyPortalIO::USkyPortalIO()
{
	bPortalReady = OpenPortalHandle();


}

bool USkyPortalIO::OpenPortalHandle() {
	//reset
	if (PortalDevice) {
		hid_close(PortalDevice);
		bPortalReady = false;
	}
	/*
		Declare two pointers to hold information about HID devices.
		"list" will point to the head of the linked list of devices,
		"attributes" will be used to iterate through the list.
	*/
	struct hid_device_info* list, * attributes;

	list = hid_enumerate(0x0, 0x0);
	// If `list` is NULL, that means no devices were found or there was an error.
	// In this case, print an error message and terminate the program.
	if (!list) {
		UE_LOG(LogHIDApi, Error, TEXT("No devices found"));
		// Get the error message from the HIDAPI
		HidError = hid_error(NULL);
		UE_LOG(LogHIDApi, Error, TEXT("%s"), *HidError);
		return false;
	}
	attributes = list;




	// Iterate through the linked list of devices
	int vendorCount = sizeof(VendorIds) / sizeof(VendorIds[0]);
	int productCount = sizeof(ProductIds) / sizeof(ProductIds[0]);

	while (attributes) {
		// Check if the devices match any of vendor_id and product_id 
		for (int i = 0; i < vendorCount; i++) {
			for (int j = 0; j < productCount; j++) {
				if (attributes->vendor_id == VendorIds[i] && attributes->product_id == ProductIds[j]) {
					UE_LOG(LogHIDApi, Display, TEXT("Portal found"));
					UE_LOG(LogHIDApi, Log, TEXT("Vendor ID: 0x%x, Product ID: 0x%x"), attributes->vendor_id, attributes->product_id);

					PortalDevice = hid_open(attributes->vendor_id, attributes->product_id, NULL);
					if (PortalDevice) {
						UE_LOG(LogHIDApi, Display, TEXT("Successful connection to Portal."));

						// Free the device list
						hid_free_enumeration(list);
						return true;
					}

					break;
				}
			}
		}

		// Move to the next device in the list
		attributes = attributes->next;
	}

	// Free the device list
	hid_free_enumeration(list);
	UE_LOG(LogHIDApi, Error, TEXT("No Portals found"));
	return false;
}

// write function need to be different on windows, as hid_write doesn't work with the way windows handle I/O
#if PLATFORM_WINDOWS

#include <Windows.h>

#define HID_CTL_CODE(id)    \
        CTL_CODE(FILE_DEVICE_KEYBOARD, (id), METHOD_NEITHER, FILE_ANY_ACCESS)
#define HID_IN_CTL_CODE(id)  \
        CTL_CODE(FILE_DEVICE_KEYBOARD, (id), METHOD_IN_DIRECT, FILE_ANY_ACCESS)
#define IOCTL_HID_SET_OUTPUT_REPORT             HID_IN_CTL_CODE(101)

struct hid_device_ {
	HANDLE device_handle;
	BOOL blocking;
	USHORT output_report_length;
	unsigned char* write_buf;
	size_t input_report_length;
	USHORT feature_report_length;
	unsigned char* feature_buf;
	wchar_t* last_error_str;
	BOOL read_pending;
	char* read_buf;
	OVERLAPPED ol;
	OVERLAPPED write_ol;
	struct hid_device_info* device_info;
};


bool USkyPortalIO::Write(FWriteBlock* pb)
{
	if (!ensure(PortalDevice)) {
		UE_LOG(LogSkyportalIO, Error, TEXT("No Portal found"));
		return false;
	}
	BOOL res;

	OVERLAPPED ol;
	memset(&ol, 0, sizeof(ol));

	DWORD bytes_returned;

	res = DeviceIoControl(PortalDevice->device_handle,
		IOCTL_HID_SET_OUTPUT_REPORT,
		(unsigned char*)pb->data, write_buf_size,
		(unsigned char*)pb->data, write_buf_size,
		&bytes_returned, &ol);
	ensureMsgf(res, TEXT("Unable to write to Portal"));
	pb->BytesTransferred = bytes_returned;
	return res;

}

#else
bool USkyPortalIO::Write(FWriteBlock* pb) {

	if (!ensure(PortalDevice)) {
		UE_LOG(LogSkyportalIO, Error, TEXT("No Portal found"));
		return;
	}

	pb->data[0] = 0; // Use report 0

	res = hid_write(PortalDevice, pb->data, write_buf_size);
	ensureMsgf(res != -1, TEXT("Unable to write to Portal, %s"), hid_error(PortalDevice));
	pb->BytesTransferred = res;
	return res;
}
#endif

void USkyPortalIO::WriteRaw(const TArray<uint8>* block)
{
	;
	int res = hid_write(PortalDevice, block->GetData(), block->GetAllocatedSize()); //return the number of byte written
	if (res == -1)
	{
		UE_LOG(LogSkyportalIO, Error, TEXT("Unable to write raw data to Portal. error:\n %s"), hid_error(PortalDevice));
		return;
	}
}

uint8* USkyPortalIO::Read()
{
	uint8* output = new uint8[0x20];
	int result = hid_read(PortalDevice, output, 0x20);

	if (result == -1) {
		UE_LOG(LogSkyportalIO, Error, TEXT("Unable to read data from Portal. Error:\n%s"), hid_error(PortalDevice));
		delete[] output;  // Prevent memory leak
		return nullptr;
	}

	if (result == 0) {  // No data was read
		UE_LOG(LogSkyportalIO, Warning, TEXT("No data read from portal"));
		delete[] output;
		return nullptr;
	}
	return output;
}

uint8* USkyPortalIO::QueryBlock(uint8 FigureIndex, uint8 BlockIndex)
{
	UE_LOG(LogSkyportalIO, Verbose, TEXT("Querying block %d - Start"), BlockIndex);
	FWriteBlock command;
	memset(command.data, 0, write_buf_size); //maybe not needed here
	command.data[1] = 'Q';
	command.data[2] = FigureIndex;  // Figure index (0x00-0x0F)
	command.data[3] = BlockIndex;
	unsigned char* output = nullptr;
	int attempt=0;
	bool write;
	do {
		attempt++;
		UE_LOG(LogSkyportalIO, Verbose, TEXT("Trying to write... attempt:%d"), attempt);
		write = Write(&command);
		if (write) {
			output = Read();
			if (output != nullptr && output != 0) {
				UE_LOG(LogSkyportalIO, Verbose, TEXT("Read success: output[0]=%d, output[1]=%d, output[2]=%d"), output[0], output[1], output[2]);
			}else{
				UE_LOG(LogSkyportalIO, Error, TEXT("Read failed, output is null"));
				return 0;
			}
		}
				
	} while (write == false && output && (output[0] != 'Q' || (output[1] % 0x10 != FigureIndex && output[1] != 0x01) || output[2] != BlockIndex) && attempt < 10);
	if (output == nullptr) {
		UE_LOG(LogSkyportalIO, Error, TEXT("Query failed after %d attempts"), attempt);
		return 0;  // Return early if the loop failed to get valid output
	}
	UE_LOG(LogSkyportalIO, Verbose, TEXT("Querying block %d - success"), BlockIndex);
	UE_LOG(LogSkyportalIO, VeryVerbose, TEXT("Data block %d = \n %s"), BlockIndex,*OutputToString(output));
	return output;
}

void USkyPortalIO::Close() {
	if (PortalDevice) {
		hid_close(PortalDevice);
	}
	bPortalReady = false;
}


/* TODO: Should not be here
bool USkyPortalIO::IsFalsePositive() const
{
	int dif = FMath::Abs(CurrentStatusData.Counter - OldStatusData.Counter);
	return (dif <= 2 || dif >= 254);

}
*/

EPortalCommand GetPortalCommandFromChar(unsigned char Char)
{
	switch (Char)
	{
	case 'A':
		return EPortalCommand::A;
	case 'C':
		return EPortalCommand::C;
	case 'J':
		return EPortalCommand::J;
	case 'L':
		return EPortalCommand::L;
	case 'M':
		return EPortalCommand::M;
	case 'Q':
		return EPortalCommand::Q;
	case 'R':
		return EPortalCommand::R;
	case 'S':
		return EPortalCommand::S;
	default:
		// Handle the case when the character doesn't match any enum
		// Return a default or invalid value, or handle the error
		UE_LOG(LogSkyportalIO, Warning, TEXT("Invalid character for Portal Command: %c"), TCHAR(Char));
		return EPortalCommand::S;  // 'S' for Status as a default
	}

}

FString DataToString(const uint8* data)
{
		FString OutString;
		for (int32 i = 0; i < 16; i++)
		{
			OutString += FString::Printf(TEXT("%02x"), data[i]);
		}
		return OutString;
}

FString OutputToString(const unsigned char* output) {
	FString OutString;
	for (int32 i = 0; i < 0x20; i++)
	{
		OutString += FString::Printf(TEXT("%02x"), output[i]);
	}
	return OutString;
}

FPortalStatusData ParsePortalStatus(const uint8* StatusResponse)
{
	FPortalStatusData result;
	result.Counter = StatusResponse[5];
	result.bIsReady = (StatusResponse[6] == 0x01);

	// Parse the figure status array (little-endian 32-bit integer)
	uint32 FigureStatusArray = 0;
	// Reading the 32-bit integer (character status array) from the buffer
	FigureStatusArray |= StatusResponse[1]; // 1st byte
	FigureStatusArray |= (StatusResponse[2] << 8);  // 2nd byte
	FigureStatusArray |= (StatusResponse[3] << 16); // 3rd byte
	FigureStatusArray |= (StatusResponse[4] << 24); // 4th byte

	bool bChangeBitsSet = false;

	TStaticArray<EFigureStatus, 16> tempArray;
	// For each of the 16 entries, extract the 2-bit status and map it to EFigureStatus
	for (int32 i = 0; i < 16; ++i)
	{
		uint8 StatusBits = (FigureStatusArray >> (i * 2)) & 0b11;  // Extract 2 bits
		EFigureStatus FigureStatus;

		switch (StatusBits)
		{
		case 0b00:
			FigureStatus = EFigureStatus::NOT_PRESENT;
			break;
		case 0b01:
			FigureStatus = EFigureStatus::PRESENT;
			break;
		case 0b11:
			FigureStatus = EFigureStatus::ADDED; // When it's a new NUID detected
			break;
		case 0b10:
			FigureStatus = EFigureStatus::REMOVED; // New NUID removed
			break;
		default:
			FigureStatus = EFigureStatus::NOT_PRESENT;  // Default case
			break;
		}

		// Add to the array of figure statuses
		//PortalStatusData.StatusArray.Insert(FigureStatus, i);
		tempArray[i] = FigureStatus;
	}
	result.StatusArray.SetNum(0);
	result.StatusArray.Append(tempArray);

	return result;
}


void DecryptAES128(uint8* OutData, const uint8* InData, const uint8* Key)
{
	UE_LOG(LogSkyportalIO, VeryVerbose, TEXT("DECRYPT -- Starting decryting data %s with %s key"), *DataToString(InData), *DataToString(Key));
	AES_KEY AesKey;
	// Set the decryption key (16 bytes for AES-128)
	AES_set_decrypt_key(Key, 128, &AesKey);

	// Perform AES-128 decryption (ECB mode)
	AES_ecb_encrypt(InData, OutData, &AesKey, AES_DECRYPT);
	UE_LOG(LogSkyportalIO, VeryVerbose, TEXT("DECRYPT -- data %s with %s key decryted"), *DataToString(InData), *DataToString(Key));
}

UFigureData* USkyPortalIO::ReadFigureBlocks(uint8 FigureIndex)
{

	UFigureData* TempFigureData = NewObject<UFigureData>();
	TempFigureData->dataError = false;  // Initialize error flag



	// Loop over all 64 blocks
	for (uint8 BlockIndex = 0; BlockIndex < FIGURE_TOTAL_BLOCKS; ++BlockIndex)
	{
		UE_LOG(LogSkyportalIO, VeryVerbose, TEXT("Reading block %d"), BlockIndex);
		
		// Query the block from the portal
		uint8* output = QueryBlock(FigureIndex, BlockIndex);
		if (output == 0) {
			--BlockIndex;
			UE_LOG(LogSkyportalIO, Error, TEXT("Query error, reprocessing..."));
			continue;
		}
			
		// Copy 16 bytes from the output, starting at the third byte
		FMemory::Memcpy(TempFigureData->data[BlockIndex], output + 3, FIGURE_BLOCK_SIZE);


		// Block 1 is sometimes a duplicate of block 0
		if (BlockIndex == 1)
		{
			if (FMemory::Memcmp(TempFigureData->data[1], TempFigureData->data[0], FIGURE_BLOCK_SIZE) == 0)
			{
				// Decrement BlockIndex to reprocess block 1 if it's a duplicate of block 0
				--BlockIndex;
				UE_LOG(LogSkyportalIO, Verbose, TEXT("Block 1 is a duplicate of block 0, reprocessing..."));
				continue;
			}
		}

		if (((BlockIndex + 1) % 4 == 0) || BlockIndex < 8)
		{
			// Direct copy from data to decryptedData for certain blocks
			FMemory::Memcpy(TempFigureData->decryptedData[BlockIndex], TempFigureData->data[BlockIndex], FIGURE_BLOCK_SIZE);
			UE_LOG(LogSkyportalIO, Verbose, TEXT("Block %d doesnt need decryption"), BlockIndex);
		}
		else {
			/***** MD5 Hash Calculation *****/
			// Prepare the hash input buffer
			UE_LOG(LogSkyportalIO, Verbose, TEXT("DECRYPT -- Starting decryting block %d"), BlockIndex);
			uint8 hashIn[0x56];
			FMemory::Memcpy(hashIn, TempFigureData->data[0], 16);
			FMemory::Memcpy(hashIn + 0x10, TempFigureData->data[1], 16);
			hashIn[0x20] = BlockIndex;
			// Append the 35-byte constant
			FMemory::Memcpy(hashIn + 0x21, HASH_CONST, sizeof(HASH_CONST));

			// Compute the MD5 hash
			FMD5 MD5;
			MD5.Update(hashIn, sizeof(hashIn));
			uint8 key[16];  // MD5 produces 16 bytes
			MD5.Final(key);

			// AES-128 ECB decryption

			FAES::FAESKey AesKey;
			FMemory::Memcpy(AesKey.Key, key, 16);

			//TODO: FAES implementation in Unreal is AES-256. We need AES-128
			//FAES::DecryptData(TempFigureData.decryptedData[BlockIndex], TempFigureData.data[BlockIndex], 16, AesKey);
			DecryptAES128(TempFigureData->decryptedData[BlockIndex], TempFigureData->data[BlockIndex], key);
			UE_LOG(LogSkyportalIO, Verbose, TEXT("DECRYPT -- block %d decrypted"), BlockIndex);
		}
		UE_LOG(LogSkyportalIO, VeryVerbose, TEXT("block %d readed with success"), BlockIndex);

	}

	// Some verifications should happen. Like if dataError is set.

	return TempFigureData;  // Return the complete figure data
}