#include <stdio.h>
#include <stdint.h>
#include <stdlib.h>
#include <string.h>
#include <strings.h>
#include <unistd.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <errno.h>

#define UPKG_MAGIC 0x9E2A83C1

// uncomment if you want full data dumps, helpful if you need to reverse engineer some unsupported format
//#define _DEBUG 1

typedef struct
{
	uint32_t magic;
	uint16_t pkgver, license;
	uint32_t flags, nnames, onames, nexports, oexports, nimports, oimports;
} upkg_header_t;

uint8_t *pkgfile;
upkg_header_t *head;
size_t fpos = 0;

uint8_t readbyte( void )
{
	uint8_t val = pkgfile[fpos];
	fpos++;
	return val;
}

uint16_t readword( void )
{
	uint16_t val = *(uint16_t*)(pkgfile+fpos);
	fpos += 2;
	return val;
}

uint32_t readdword( void )
{
	uint32_t val = *(uint32_t*)(pkgfile+fpos);
	fpos += 4;
	return val;
}

float readfloat( void )
{
	float val = *(float*)(pkgfile+fpos);
	fpos += 4;
	return val;
}

#define READSTRUCT(x,y) {memcpy(&x,pkgfile+fpos,sizeof(y));fpos+=sizeof(y);}

// reads a compact index value
int32_t readindex( void )
{
	uint8_t byte[5] = {0};
	byte[0] = readbyte();
	if ( !byte[0] ) return 0;
	if ( byte[0]&0x40 )
	{
		for ( int i=1; i<5; i++ )
		{
			byte[i] = readbyte();
			if ( !(byte[i]&0x80) ) break;
		}
	}
	int32_t tf = byte[0]&0x3f;
	tf |= (int32_t)(byte[1]&0x7f)<<6;
	tf |= (int32_t)(byte[2]&0x7f)<<13;
	tf |= (int32_t)(byte[3]&0x7f)<<20;
	tf |= (int32_t)(byte[4]&0x7f)<<27;
	if ( byte[0]&0x80 ) tf *= -1;
	return tf;
}

// reads a name table entry
size_t readname( int *olen )
{
	size_t pos = fpos;
	if ( head->pkgver >= 64 )
	{
		int32_t len = readindex();
		pos = fpos;
		if ( olen ) *olen = len;
		if ( len <= 0 ) return pos;
		fpos += len;
	}
	else
	{
		int c, p = 0;
		while ( (c = readbyte()) ) p++;
		if ( olen ) *olen = p;
	}
	fpos += 4;
	return pos;
}

size_t getname( int index, int *olen )
{
	size_t prev = fpos;
	fpos = head->onames;
	size_t npos = 0;
	for ( int i=0; i<=index; i++ )
		npos = readname(olen);
	fpos = prev;
	return npos;
}

// checks if a name exists
int hasname( const char *needle )
{
	if ( !needle ) return 0;
	size_t prev = fpos;
	fpos = head->onames;
	int found = 0;
	int nlen = strlen(needle);
	for ( uint32_t i=0; i<head->nnames; i++ )
	{
		int32_t len = 0;
		if ( head->pkgver >= 64 )
		{
			len = readindex();
			if ( len <= 0 ) continue;
		}
		int c = 0, p = 0, match = 1;
		while ( (c = readbyte()) )
		{
			if ( (p >= nlen) || (needle[p] != c) ) match = 0;
			p++;
			if ( len && (p > len) ) break;
		}
		if ( match )
		{
			found = 1;
			break;
		}
		fpos += 4;
	}
	fpos = prev;
	return found;
}

// read import table entry and return index of its object name
int32_t readimport( void )
{
	readindex();
	readindex();
	if ( head->pkgver >= 55 ) fpos += 4;
	else readindex();
	return readindex();
}

int32_t getimport( int index )
{
	size_t prev = fpos;
	fpos = head->oimports;
	int32_t iname = 0;
	for ( int i=0; i<=index; i++ )
		iname = readimport();
	fpos = prev;
	return iname;
}

// fully read import table entry
void readimport2( int32_t *cpkg, int32_t *cname, int32_t *pkg, int32_t *name )
{
	*cpkg = readindex();
	*cname = readindex();
	if ( head->pkgver >= 55 ) *pkg = readdword();
	else *pkg = readindex();
	*name = readindex();
}

void getimport2( int index, int32_t *cpkg, int32_t *cname, int32_t *pkg,
	int32_t *name )
{
	size_t prev = fpos;
	fpos = head->oimports;
	for ( int i=0; i<=index; i++ )
		readimport2(cpkg,cname,pkg,name);
	fpos = prev;
}

// read export table entry
void readexport( int32_t *class, int32_t *ofs, int32_t *siz, int32_t *name )
{
	*class = readindex();
	readindex();
	if ( head->pkgver >= 55 ) fpos += 4;
	*name = readindex();
	fpos += 4;
	*siz = readindex();
	if ( *siz > 0 ) *ofs = readindex();
}

void getexport( int index, int32_t *class, int32_t *ofs, int32_t *siz,
	int32_t *name )
{
	size_t prev = fpos;
	fpos = head->oexports;
	for ( int i=0; i<=index; i++ )
		readexport(class,ofs,siz,name);
	fpos = prev;
}

// fully read export table entry
void readexport2( int32_t *class, int32_t *super, int32_t *pkg, int32_t *name,
	uint32_t *flags, int32_t *siz, int32_t *ofs )
{
	*class = readindex();
	*super = readindex();
	if ( head->pkgver >= 55 ) *pkg = readdword();
	else *pkg = 0;
	*name = readindex();
	*flags = readdword();
	*siz = readindex();
	if ( *siz > 0 ) *ofs = readindex();
}

void getexport2( int index, int32_t *class, int32_t *super, int32_t *pkg,
	int32_t *name, uint32_t *flags, int32_t *siz, int32_t *ofs )
{
	size_t prev = fpos;
	fpos = head->oexports;
	for ( int i=0; i<=index; i++ )
		readexport2(class,super,pkg,name,flags,siz,ofs);
	fpos = prev;
}

typedef struct
{
	uint8_t r, g, b, x;
} __attribute__((packed)) color_t;

typedef struct
{
	uint8_t r, g, b;
} __attribute__((packed)) pal8_t;

const char sparktypes[29][32] =
{
	"Burn", "Sparkle", "Pulse", "Signal", "Blaze",
	"OzHasSpoken", "Cone", "BlazeRight", "BlazeLeft",
	"Cylinder", "Cylinder3D", "Lissajous", "Jugglers",
	"Emit", "Fountain", "Flocks", "Eels", "Organic",
	"WanderOrganic", "RandomCloud", "CustomCloud",
	"LocalCloud", "Stars", "LineLighting", "RampLighting",
	"SphereLighting", "Wheel", "Gametes", "Sprinkler"
};

void savepalette( uint32_t pal, int32_t *palnamelen, char **palname )
{
	size_t prev = fpos;
	int32_t class, ofs, siz, name;
	getexport(pal-1,&class,&ofs,&siz,&name);
	*palname = (char*)(pkgfile+getname(name,palnamelen));
	// begin reading data
	fpos = ofs;
	if ( head->pkgver < 45 ) fpos += 4;
	if ( head->pkgver < 55 ) fpos += 16;
	if ( head->pkgver <= 44 ) fpos -= 6;	// ???
	if ( head->pkgver == 45 ) fpos -= 2;	// ???
	if ( head->pkgver <= 35 ) fpos += 8;	// ???
	// process properties
	int32_t prop = readindex();
	if ( (uint32_t)prop >= head->nnames )
	{
		printf("Unknown property %d, skipping\n",prop);
		fpos = prev;
		return;
	}
	int32_t l = 0;
	char *pname = (char*)(pkgfile+getname(prop,&l));
retrypal:
	if ( strncasecmp(pname,"none",l) )
	{
		uint8_t info = readbyte();
		//int array = info&0x80;
		int type = info&0xf;
		int psiz = (info>>4)&0x7;
		switch ( psiz )
		{
		case 0:
			psiz = 1;
			break;
		case 1:
			psiz = 2;
			break;
		case 2:
			psiz = 4;
			break;
		case 3:
			psiz = 12;
			break;
		case 4:
			psiz = 16;
			break;
		case 5:
			psiz = readbyte();
			break;
		case 6:
			psiz = readword();
			break;
		case 7:
			psiz = readdword();
			break;
		}
		if ( type == 10 )
			readindex();	// skip struct name
		fpos += psiz;
		prop = readindex();
		pname = (char*)(pkgfile+getname(prop,&l));
		goto retrypal;
	}
	int32_t grouplen = 0;
	char *group = 0;
	if ( head->pkgver < 55 )
	{
		// group
		fpos++;
		int32_t grp = readindex();
		group = (char*)(pkgfile+getname(grp,&grouplen));
	}
	// number of colors
	int32_t num = readindex();
	char fname[256];
	if ( group && strncmp(group,"None",grouplen) )
	{
		snprintf(fname,256,"Palettes/%.*s",grouplen,group);
		mkdir(fname,0775);
		snprintf(fname,256,"Palettes/%.*s/%.*s.pal",grouplen,group,
			*palnamelen,*palname);
	}
	else snprintf(fname,256,"Palettes/%.*s.pal",*palnamelen,*palname);
	color_t *paldat = (color_t*)(pkgfile+fpos);
	pal8_t *palconv = calloc(sizeof(pal8_t),num);
	for ( int i=0; i<num; i++ )
	{
		palconv[i].r = paldat[i].r;
		palconv[i].g = paldat[i].g;
		palconv[i].b = paldat[i].b;
		// discard X
	}
	fpos = prev;
	int fd = open(fname,O_WRONLY|O_CREAT|O_TRUNC,0644);
	if ( fd == -1 )
	{
		printf("Failed to open file %s: %s\n",fname,strerror(errno));
		free(palconv);
		return;
	}
	int w = write(fd,palconv,sizeof(pal8_t)*num);
	if ( w == -1 )
	{
		close(fd);
		free(palconv);
		printf("Write failed for file %s: %s\n",fname,strerror(errno));
		return;
	}
	close(fd);
	free(palconv);
}

void savetexture( int32_t namelen, char *name, uint32_t pal, int brising,
	uint8_t renderheat, uint32_t sparklimit, int version, int32_t grouplen,
	char *group )
{
	if ( version < 55 )
	{
		// group
		fpos++;
		int32_t grp = readindex();
		group = (char*)(pkgfile+getname(grp,&grouplen));
	}
	uint8_t mipcnt = readbyte();
	uint32_t ofs = 0;
	if ( version >= 63 ) ofs = readdword();
	uint32_t datasiz = readindex();
	if ( version >= 63 ) fpos = ofs;
	else fpos += datasiz;
	uint32_t w = readdword();
	uint32_t h = readdword();
	readbyte();	// ubits
	readbyte();	// vbits
	// skip all the other mipmaps, we only need the dimensions of mip zero
	for ( int i=1; i<mipcnt; i++ )
	{
		if ( version >= 63 ) ofs = readdword();
		datasiz = readindex();
		if ( version >= 63 ) fpos = ofs;
		else fpos += datasiz;
		readdword();
		readdword();
		readbyte();
		readbyte();
	}
	char fname[256];
	if ( group && strncmp(group,"None",grouplen) )
	{
		snprintf(fname,256,"Textures/%.*s",grouplen,group);
		mkdir(fname,0775);
		snprintf(fname,256,"Textures/%.*s/%.*s.ftx",grouplen,group,
			namelen,name);
	}
	else snprintf(fname,256,"Textures/%.*s.ftx",namelen,name);
	// we're using stdio for this one, it's cheaper
	FILE *f = fopen(fname,"wb");
	if ( !f )
	{
		printf("Failed to open file %s: %s\n",fname,strerror(errno));
		return;
	}
	// texture name, width and height
	fprintf(f,"%.*s %u %u\n",namelen,name,w,h);
	// palette name
	int32_t palnamelen = 0;
	char *palname = 0;
	savepalette(pal,&palnamelen,&palname);
	fprintf(f,"%.*s\n",palnamelen,palname);
	// bRising, RenderHeat and SparkLimit
	fprintf(f,"%d %hhu %u\n",!!brising,renderheat,sparklimit);
	// individual sparks (type as string, heat, x, y, args[4])
	int32_t nsparks = readindex();	// TArray<FSpark> size
	for ( int32_t i=0; i<nsparks; i++ )
	{
		uint8_t type, heat, x, y, args[4];
		type = readbyte();
		heat = readbyte();
		x = readbyte();
		y = readbyte();
		args[0] = readbyte();
		args[1] = readbyte();
		args[2] = readbyte();
		args[3] = readbyte();
		if ( type >= 29 ) continue;	// discard unknown types
		fprintf(f,"%s %hhu %hhu %hhu %hhu %hhu %hhu %hhu\n",
			sparktypes[type],heat,x,y,
			args[0],args[1],args[2],args[3]);
	}
	fclose(f);
}

int main( int argc, char **argv )
{
	if ( argc < 2 )
	{
		printf("Usage: uftxextract <archive>\n");
		return 1;
	}
	int fd = open(argv[1],O_RDONLY);
	if ( fd == -1 )
	{
		printf("Failed to open file %s: %s\n",argv[1],strerror(errno));
		return 1;
	}
	struct stat st;
	fstat(fd,&st);
	pkgfile = malloc(st.st_size);
	memset(pkgfile,0,st.st_size);
	head = (upkg_header_t*)pkgfile;
	int r = 0;
	do
	{
		r = read(fd,pkgfile+fpos,131072);
		if ( r == -1 )
		{
			close(fd);
			free(pkgfile);
			printf("Read failed for file %s: %s\n",argv[1],
				strerror(errno));
			return 4;
		}
		fpos += r;
	}
	while ( r > 0 );
	close(fd);
	fpos = 0;
	if ( head->magic != UPKG_MAGIC )
	{
		printf("File %s is not a valid unreal package!\n",argv[1]);
		free(pkgfile);
		return 2;
	}
	// loop through exports and search for textures
	fpos = head->oexports;
	for ( uint32_t i=0; i<head->nexports; i++ )
	{
		int32_t class, super, pkg, name, siz, ofs;
		uint32_t flags;
		readexport2(&class,&super,&pkg,&name,&flags,&siz,&ofs);
		if ( (siz <= 0) || (class >= 0) ) continue;
		// get the class name
		class = -class-1;
		if ( (uint32_t)class > head->nimports ) continue;
		int32_t l = 0;
		char *n = (char*)(pkgfile+getname(getimport(class),&l));
		int istex = !strncasecmp(n,"FireTexture",l);
		if ( !istex ) continue;
		mkdir("Textures",0775);
		mkdir("Palettes",0775);
		// get the highest group name (must be an export)
		char *pkgn = 0;
		int32_t pkgl = 0;
		while ( pkg > 0 )
		{
			int32_t pclass, psuper, ppkg, pname, psiz, pofs;
			uint32_t pflags;
			getexport2(pkg-1,&pclass,&psuper,&ppkg,&pname,&pflags,&psiz,&pofs);
			pkgn = (char*)(pkgfile+getname(pname,&pkgl));
			pkg = ppkg;
		}
		char *tex = (char*)(pkgfile+getname(name,&l));
		if ( pkgn && strncmp(pkgn,"None",pkgl) )
			printf("FireTexture found: %.*s.%.*s\n",pkgl,pkgn,l,tex);
		else printf("FireTexture found: %.*s\n",l,tex);
		int32_t texl = l;
#ifdef _DEBUG
		char fname[256] = {0};
		snprintf(fname,256,"%.*s.object",texl,tex);
		printf(" Dumping full object data to %s\n",fname);
		FILE *f = fopen(fname,"wb");
		fwrite(pkgfile+ofs,siz,1,f);
		fclose(f);
		continue;
#endif
		// begin reading data
		size_t prev = fpos;
		fpos = ofs;
		if ( head->pkgver < 45 ) fpos += 4;
		if ( head->pkgver < 55 ) fpos += 16;
		if ( head->pkgver <= 44 ) fpos -= 6;	// ???
		if ( head->pkgver == 45 ) fpos -= 2;	// ???
		if ( head->pkgver == 41 ) fpos += 2;	// ???
		if ( head->pkgver <= 35 ) fpos += 8;	// ???
		// process properties
		int32_t prop = readindex();
		if ( (uint32_t)prop >= head->nnames )
		{
			printf("Unknown property %d, skipping\n",prop);
			fpos = prev;
			continue;
		}
		int32_t pal = 0;	// we need to fetch this
		char *pname = (char*)(pkgfile+getname(prop,&l));
		uint8_t rising = 0;	// as well as these
		uint8_t heat = 200;
		uint32_t limit = 1024;
retry:
		if ( strncasecmp(pname,"None",l) )
		{
			uint8_t info = readbyte();
			int array = info&0x80;
			int type = info&0xf;
			int psiz = (info>>4)&0x7;
			//int32_t tl = 0;
			//char *sname = 0;
			if ( type == 10 )
			{
				//int32_t sn;
				/*sn = */readindex();
				//sname = (char*)(pkgfile+getname(sn,&tl));
			}
			switch ( psiz )
			{
			case 0:
				psiz = 1;
				break;
			case 1:
				psiz = 2;
				break;
			case 2:
				psiz = 4;
				break;
			case 3:
				psiz = 12;
				break;
			case 4:
				psiz = 16;
				break;
			case 5:
				psiz = readbyte();
				break;
			case 6:
				psiz = readword();
				break;
			case 7:
				psiz = readdword();
				break;
			}
			//printf(" prop %.*s (%u, %u, %u, %u)\n",l,pname,array,type,(info>>4)&7,psiz);
			//if ( tl ) printf("  struct: %.*s\n",tl,sname);
			if ( array && (type != 3) )
			{
				/*int idx = */readindex();
				//printf(" index: %d\n",idx);
			}
			if ( !strncasecmp(pname,"Palette",l) )
				pal = readindex();
			else if ( !strncasecmp(pname,"bRising",l) )
				rising = (info>>7)&1;
			else if ( !strncasecmp(pname,"RenderHeat",l) )
				heat = readbyte();
			else if ( !strncasecmp(pname,"SparksLimit",l) )
				limit = readdword();
			else fpos += psiz;
			prop = readindex();
			pname = (char*)(pkgfile+getname(prop,&l));
			goto retry;
		}
		if ( pal )
			savetexture(texl,tex,pal,rising,heat,limit,head->pkgver,pkgl,pkgn);
		fpos = prev;
	}
	free(pkgfile);
	return 0;
}