#include <vector>
#include <string>
#include <fstream>
#include <iostream>

#include <cstring>
#include <cstdio>

#include <sys/stat.h>
#include <dirent.h>

#include <lz4frame.h>

#define PATH_PART_LEN 32

/* Resource file is in simple structure:
 * 4 bytes: ID
 * 2 bytes: Data size
 * n bytes: Data
 *
 * Subchunks come after data, so you have to know how chunk is strcutured
 *
 * All numbers are in native endian
 *
 * DIRE chunk comes with following data:
 *   PATH_PART_LEN bytes name
 * and nests FILE chunks
 *
 * FILE chunk comes with following data:
 *   2             bytes flag
 *   4             bytes data location
 *   4             bytes data size
 *   PATH_PART_LEN bytes name
 */

struct Data {
	unsigned int size;
	unsigned char* data;
};

#define FILE_FLAG_RAW (1 << 0)

unsigned int loc = 0;
std::vector<struct Data> datalist;

class Chunk;

class Chunk {
	void init(std::string name){
		this->data = nullptr;
		memset(this->id, 0, 4);
		strcpy(this->id, name.c_str());
		this->name = "";
		this->fullpath = "";
		this->keepRaw = false;
	}

public:
	char id[4];
	unsigned char* data;
	unsigned int size;
	Chunk* parent;
	std::vector<Chunk*> subchunks;
	std::string name;
	std::string fullpath;
	bool keepRaw;

	Chunk(std::string name){
		this->init(name);

		this->parent = nullptr;
	}

	Chunk(Chunk* parent, std::string name){
		this->init(name);

		this->parent = parent;
		this->parent->subchunks.push_back(this);
	}

	~Chunk(){
		if(this->data != nullptr) delete[] this->data;

		for(int i = 0; i < this->subchunks.size(); i++){
			delete this->subchunks[i];
		}
	}

	unsigned int totalSize(){
		unsigned int sz = 0;

		for(int i = 0; i < this->subchunks.size(); i++){
			sz = this->subchunks[i]->totalSize();
		}

		return 4 + 4 + this->size + sz;
	}

	void write(std::ofstream* ofs){
		unsigned int sz = 0;

		if(memcmp(this->id, "FILE", 4) == 0){
			unsigned short* flag;
			unsigned int* location;
			unsigned int* size;
			char* name;
			std::ifstream ifs(this->fullpath, std::ios::binary);

			unsigned int cbound;
			unsigned char* fdata;
			unsigned char* cdata;
			unsigned int fsize;
			unsigned int csize;
			struct Data dc;

			ifs.seekg(0, std::ios_base::end);
			fsize = ifs.tellg();
			ifs.seekg(0, std::ios_base::beg);

			fdata = new unsigned char[fsize];
			ifs.read((char*)fdata, fsize);

			if(this->keepRaw){
				cdata = fdata;
				csize = fsize;
			}else{
				cbound = LZ4F_compressFrameBound(fsize, nullptr);
				cdata = new unsigned char[cbound];

				csize = LZ4F_compressFrame(cdata, cbound, fdata, fsize, nullptr);

				delete[] fdata;
			}

			this->size = 2 + 4 + 4 + PATH_PART_LEN;
			this->data = new unsigned char[this->size];

			flag = (unsigned short*)&this->data[0];
			location = (unsigned int*)&this->data[2];
			size = (unsigned int*)&this->data[2 + 4];
			name = (char*)&this->data[2 + 4 + 4];

			*flag = (this->keepRaw ? FILE_FLAG_RAW : 0);
			*location = loc;
			*size = csize;

			memset(name, 0, PATH_PART_LEN);
			memcpy(name, this->name.c_str(), this->name.size());

			dc.data = cdata;
			dc.size = csize;
			datalist.push_back(dc);

			loc += csize;
		}else if(memcmp(this->id, "DIRE", 4) == 0){
			char* name;

			this->size = PATH_PART_LEN;
			this->data = new unsigned char[this->size];

			name = (char*)&this->data[0];
			
			memset(name, 0, PATH_PART_LEN);
			memcpy(name, this->name.c_str(), this->name.size());
		}
		
		sz = this->totalSize() - 4 - 4;

		ofs->write((char*)this->id, 4);
		ofs->write((char*)&sz, 4);
		if(this->size > 0) ofs->write((char*)this->data, this->size);
		for(int i = 0; i < this->subchunks.size(); i++){
			this->subchunks[i]->write(ofs);
		}
	}
};

static void readDir(Chunk* parent, std::string dirname){
	DIR* dir = opendir(dirname.c_str());
	struct dirent* d;

	if(dir != nullptr){
		while((d = readdir(dir)) != nullptr){
			struct stat s;
			std::string dn(d->d_name);
			std::string p = dirname + "/" + dn;
			Chunk* c;
			
			if(dn == ".." || dn == ".") continue;

			stat(p.c_str(), &s);

			if(S_ISDIR(s.st_mode)){
				c = new Chunk(parent, "DIRE");
				c->name = dn;

				readDir(c, p);
			}else{
				int pos = dn.find_last_of('.');
				std::string ext = (pos == std::string::npos) ? "" : dn.substr(pos, dn.size() - pos);

				c = new Chunk(parent, "FILE");
				c->name = dn;
				c->fullpath = p;

				if(ext == ".png"){
					c->keepRaw = true;
				}
			}
		}
		closedir(dir);
	}
}

int main(int argc, char** argv){
	Chunk* root;
	Chunk* rootdir;
	std::ofstream* ofs;

	if(argc < 3){
		std::cerr << "Usage: " << argv[0] << " [output] [input...]" << std::endl;
		return 1;
	}
	
	root = new Chunk("FBPK");

	rootdir = new Chunk(root, "DIRE");
	rootdir->name = "";

	for(int i = 2; i < argc; i++){
		readDir(rootdir, argv[i]);
	}

	ofs = new std::ofstream(argv[1], std::ios::binary);
	root->write(ofs);
	for(int i = 0; i < datalist.size(); i++){
		ofs->write((char*)datalist[i].data, datalist[i].size);
	}
	ofs->close();
}