#include "resourceexplorer.h"
#include "ui_resourceexplorer.h"

#include <QFile>
#include <QFileInfo>
#include <QDateTime>
#include <QDebug>
#include <QFileDialog>

#include <endian.h>
#include <list>

#include "injector.h"

#pragma pack (push, 1)


struct ResourceLibraryEntry {
	uint32_t nameOffset;
	uint16_t flags; // use ResourceFlags
	union Content {
		struct Directory {
			uint32_t childCount;
			uint32_t childOffset;
		} directory;

		struct File {
			uint16_t country;
			uint16_t language;
			uint32_t dataOffset;
		} file;
	} content;
};

struct ResourceLibraryEntryV2 : public ResourceLibraryEntry {
	uint64_t lastModTimeMS;
};

struct ResouceNameEntry {
	uint16_t length;
	uint32_t hash;
	// string content
	char str[0];
};

struct ResouceDataEntry {
	uint32_t length;
	char data[0];
};

#pragma pack (pop)

ResourceExplorer::ResourceExplorer( QWidget* parent ) : QMainWindow( parent ), ui( new Ui::ResourceExplorer ) {
	this->ui->setupUi(this);

	this->rebuild();
}

ResourceExplorer::~ResourceExplorer() {
	delete this->ui;
}

void ResourceExplorer::extractButtonPressed() {
	QTreeWidgetItem* item = this->ui->treeWidget->currentItem();
	if ( !item ) return;

	ResourceNode* node = (ResourceNode*) item->data(0, Qt::UserRole).toULongLong();
	if ( !node ) return;

	QString saveFile = QFileDialog::getSaveFileName(this, "Extract File");

	if ( saveFile.isEmpty() ) return;

	QFile fileOut( saveFile );
	if( !fileOut.open( QIODevice::WriteOnly | QIODevice::Truncate )) return;

	QByteArray data = QByteArray::fromRawData( (const char*) node->start, node->size );

	if ( node->isCompressed() ) {
		data = qUncompress(data);
	}

	fileOut.write(data);
}

void ResourceExplorer::itemSelected(QTreeWidgetItem* newItem) {
	if ( !newItem ) return;

	ResourceNode* node = (ResourceNode*) newItem->data(0, Qt::UserRole).toULongLong();

	if ( node ) {
		if ( !node->isDir() ) {
			this->ui->extractButton->setEnabled( true );
			return;
		}
	}
	this->ui->extractButton->setEnabled( false );
}

ResourceExplorer::ResourceNode::~ResourceNode() {
	for ( ResourceNode* rn : this->children ) {
		delete rn;
	}
}

bool ResourceExplorer::ResourceNode::isDir() const {
	return this->flags & ResourceFlags::Directory;
}

bool ResourceExplorer::ResourceNode::isCompressed() const {
	return this->flags & ResourceFlags::Compressed;
}

uint64_t ResourceExplorer::ResourceNode::getByteSize() const {
	if ( this->isDir() ) {
		return std::accumulate( children.begin(), children.end(), 0, [](uint64_t v, const ResourceNode* b) { return v + b->getByteSize(); } );
	}
	return size;
}

void ResourceExplorer::rebuild() {
	this->ui->treeWidget->clear();
	this->ui->extractButton->setEnabled(false);

	if ( !registeredResources || registeredResources->empty() ) {
		QTreeWidgetItem* resItem = new QTreeWidgetItem();
		resItem->setText( 0, "No Resources Captured" );
		this->ui->treeWidget->addTopLevelItem( resItem );
		return;
	}

	std::vector<MemoryMap> memMaps = readMemoryMaps();

	for ( auto it : *registeredResources ) {
		// find correct memory map
		auto memMapIt = std::find_if(memMaps.cbegin(), memMaps.cend(), [c = it.second](const MemoryMap& mm) { return mm.begin < c->resourceStruct && mm.end > c->resourceStruct; } );
		QString fileName( "<file unknown>" );
		QString filePath( "" );
		if ( memMapIt != memMaps.cend() ) {
			fileName = QFileInfo( memMapIt->path ).fileName();
			filePath = memMapIt->path;
		}

		QTreeWidgetItem* resItem = new QTreeWidgetItem( );

		resItem->setText( 0, fileName );
		resItem->setToolTip( 0, filePath );

		std::shared_ptr<RegisteredResource> r = it.second;
		std::vector<ResourceNode*> resources = parseResource(r->version, r->resourceStruct, r->resourceName, r->resourceData);

		for ( ResourceNode* rnode : resources ) {
			QTreeWidgetItem* item = this->resourceNodeToItem( rnode, "" );
			resItem->addChild( item );
		}

		this->ui->treeWidget->addTopLevelItem( resItem );
	}
}

std::vector<ResourceExplorer::MemoryMap> ResourceExplorer::readMemoryMaps() {
	QFile maps("/proc/self/maps");

	if ( !maps.open( QIODevice::ReadOnly | QIODevice::Text ) ) {
		qWarning() << "failed to read MemoryMaps";
		return {};
	}

	// QTextStream stream(&maps);
	std::vector<MemoryMap> out;
	out.reserve( maps.size() / 100 ); // approximation of row count
	do {
		QString line = QString::fromUtf8( maps.readLine() );
		if( line.isEmpty() ) continue;

		QVector<QStringRef> parts = line.splitRef(' ', QString::SplitBehavior::SkipEmptyParts );
		if ( parts.size() != 6) {
			continue;
		}

		QVector<QStringRef> borders = parts.at(0).split('-');
		if ( borders.size() != 2) continue;

		bool ok = true;
		const unsigned char* begin = (const unsigned char*) borders.at( 0 ).toULongLong( &ok, 16);
		if ( !ok ) continue;
		const unsigned char* end   = (const unsigned char*) borders.at( 1 ).toULongLong( &ok, 16);
		if ( !ok ) continue;

		if ( begin > end ) std::swap( begin, end );

		out.push_back( MemoryMap{ begin, end, parts.at( 5 ).toString().simplified() } );
	} while( !maps.atEnd() );

	return out;
}

// reads an UTF-16 String into a QString
// the problem is that the input string is in bigendian.
// this function makes a copy and swaps every two bytes if required
static QString readUTF16String(const char16_t* ptr, uint32_t size) {
	char16_t* newptr = new char16_t[size];

	for ( uint32_t i = 0; i < size; ++i ) {
		newptr[i] = be16toh(ptr[i]);
	}

	QString out = QString::fromUtf16(newptr, size);
	delete[] newptr;
	return out;
}

std::vector<ResourceExplorer::ResourceNode*> ResourceExplorer::parseResource(int version, ResourceExplorer::data libraryIndex, ResourceExplorer::data names, ResourceExplorer::data data ) {
	std::vector<ResourceNode*> out;

	if ( version > 2 ) {
		qWarning() << "resource has unknown version: " << version;
		return out;
	}

	const uint32_t libraryOffset = ( version == 1 ) ? sizeof(ResourceLibraryEntry) : sizeof(ResourceLibraryEntryV2);

	using task_t = std::pair<uint32_t, std::vector<ResourceNode*>*>;
	std::list<task_t> pendingOffsets;
	pendingOffsets.push_back( {0, &out} );
	while ( pendingOffsets.size() > 0) {
		task_t task = pendingOffsets.front();
		uint32_t offset = task.first;
		pendingOffsets.pop_front();

		const ResourceLibraryEntryV2* entry = (const ResourceLibraryEntryV2*) &libraryIndex[ offset * libraryOffset ];

		ResourceNode* newNode = new ResourceNode();
		newNode->flags = (ResourceFlags) be16toh( entry->flags );


		if( offset == 0) {
			newNode->name = ":";
			newNode->hash = 0;
		} else {
			uint32_t nameOffset = be32toh(entry->nameOffset);
			uint16_t nameLen = be16toh(*(uint16_t*) (names + nameOffset));
			newNode->name = readUTF16String((const char16_t*) (names + nameOffset + 6), nameLen);
			newNode->hash = *((uint16_t*) ( names + nameOffset + 2));
		}

		if ( newNode->flags & ResourceFlags::Directory ) {
			// directory
			newNode->size = 0;
			newNode->country = 0;
			newNode->language = 0;

			uint32_t childCount = be32toh( entry->content.directory.childCount );
			uint32_t childOffset = be32toh( entry->content.directory.childOffset );
			for ( uint32_t i = 0; i < childCount; ++i ) {
				pendingOffsets.push_back( { i + childOffset, &(newNode->children) } );
			}
		} else {
			// file
			newNode->country = be16toh(entry->content.file.country);
			newNode->language = be16toh(entry->content.file.language);
			uint32_t dataOffset =  be32toh(entry->content.file.dataOffset);
			const unsigned char* dataBegin = data + dataOffset;
			newNode->size = be32toh(*((const uint32_t*) dataBegin));
			newNode->start = dataBegin + 4;
		}

		if ( version >= 2 ) {
			newNode->lastModTimeMS = be64toh( entry->lastModTimeMS );
		}

		task.second->push_back( newNode );
	}

	return out;
}

QTreeWidgetItem* ResourceExplorer::resourceNodeToItem( ResourceNode* node, const QString& path ) {
	if ( ! node ) return nullptr;

	QTreeWidgetItem* item = new QTreeWidgetItem();
	const QString ownPath = path + node->name;

	item->setText( 0, node->isDir() ? (node->name + "/" ) : node->name);

	if ( node->isDir() ) {
		std::size_t size = node->children.size();
		item->setText(1, QString::number( size ));
		uint64_t filesize = node->getByteSize();
		item->setToolTip(1, QString("TotalSize: %1 (%2 Bytes)").arg(approxFileSize(filesize)).arg( filesize ));
		item->setToolTip(0, ownPath);
	} else {
		std::size_t size = node->size;
		item->setText( 1, approxFileSize( size ) );
		item->setToolTip( 1, QString::number( size ) );
		QString tooltip = ownPath;
		QLocale::Language lang = (QLocale::Language) node->language;
		QLocale::Country country = (QLocale::Country) node->country;
		QLocale fileLocale(lang, country);
		tooltip += (node->isCompressed() ? "\nCompressed" : "\nUncompressed")
			+ ((fileLocale != QLocale()) ? ("\nLocale: " + fileLocale.name() ) : "" );
		item->setToolTip(0, tooltip);
	}

	if( node->lastModTimeMS ) {
		QDateTime lastModTime = QDateTime::fromMSecsSinceEpoch( node->lastModTimeMS );
		item->setText( 2, lastModTime.toString( "yyyy-MM-dd HH:mm:ss" ) );
	} else {
		item->setText( 2, "-" );
	}

	for ( ResourceNode* subnode : node->children ) {
		QTreeWidgetItem* sub = resourceNodeToItem( subnode, ownPath + "/" );
		item->addChild( sub );
	}

	QVariant ptrVariant( (qulonglong) node );
	item->setData(0, Qt::UserRole, ptrVariant);

	return item;
}


QString ResourceExplorer::approxFileSize( uint64_t size ) {
	return QLocale().formattedDataSize( size );
}