#include "Archive.h"
#include "common/IFileStream.h"
#include "zlib/zlib.h"
#include "dds.h"

Archive::Archive()
	:m_useATIFourCC(false)
{
	//
}

Archive::~Archive()
{
	//
}

bool Archive::Read(IDataStream * src)
{
	m_src = src;

	src->ReadBuf(&m_hdr, sizeof(m_hdr));

	_MESSAGE("version = %08X", m_hdr.version);
	_MESSAGE("numFiles = %08X", m_hdr.numFiles);

	if(m_hdr.IsGeneral())
	{
		if(!Read_General())
			return false;
	}
	else if(m_hdr.IsDX10())
	{
		if(!Read_DX10())
			return false;
	}
	else
	{
		_ERROR("unhandled format: %.4s", m_hdr.type);
		return false;
	}

	UInt64 headerEnd = m_src->GetOffset();

	_MESSAGE("@%016I64X", headerEnd);

	if(!Read_Nametable())
		return false;

	return true;
}

void Archive::Extract(const char * dstRoot)
{
	if(m_hdr.IsGeneral())
		Extract_General(dstRoot);
	else if(m_hdr.IsDX10())
		Extract_DX10(dstRoot);
	else
		_ERROR("unhandled format: %.4s", m_hdr.type);
}

bool Archive::Read_General()
{
	m_files.resize(m_hdr.numFiles);
	if(m_hdr.numFiles)
		m_src->ReadBuf(&m_files[0], sizeof(FileEntry) * m_hdr.numFiles);

	for(UInt32 i = 0; i < m_files.size(); i++)
	{
		FileEntry * file = &m_files[i];

		_DMESSAGE("%08X: %08X %.4s %08X %08X %016I64X %08X %08X %08X",
			i,
			file->unk00, file->ext, file->unk08, file->unk0C,
			file->offset, file->packedLen, file->unpackedLen, file->unk20);
	}

	return true;
}

bool Archive::Read_DX10()
{
	m_textures.resize(m_hdr.numFiles);

	for(UInt32 i = 0; i < m_textures.size(); i++)
	{
		Texture * tex = &m_textures[i];

		m_src->ReadBuf(&tex->hdr, sizeof(tex->hdr));

		tex->chunks.resize(tex->hdr.numChunks);
		if(tex->hdr.numChunks)
			m_src->ReadBuf(&tex->chunks[0], sizeof(DX10Chunk) * tex->hdr.numChunks);
	}

	return true;
}

bool Archive::Read_Nametable()
{
	// name table
	m_src->SetOffset(m_hdr.nameTableOffset);

	char * strBuf = new char[0x10000];

	UInt32 idx = 0;

	while(m_src->GetRemain() >= 2)
	{
		UInt32 len = m_src->Read16();

		m_src->ReadBuf(strBuf, len);
		strBuf[len] = 0;

		m_names.push_back(strBuf);

		_DMESSAGE("%08X %s", idx, strBuf);
		idx++;
	}

	delete [] strBuf;

	return true;
}

void Archive::Extract_General(const char * dstRoot)
{
	ASSERT(m_files.size() == m_names.size());

	for(UInt32 i = 0; i < m_files.size(); i++)
	{
		FileEntry * file = &m_files[i];

		std::string dstPath = dstRoot;
		dstPath += "\\";
		dstPath += m_names[i];

		IFileStream::MakeAllDirs(dstPath.c_str());

		IFileStream dst;
		if(dst.Create(dstPath.c_str()))
		{
			m_src->SetOffset(file->offset);

			if(file->packedLen && (file->unpackedLen != file->packedLen))
			{
				UInt32 unpackedLen = file->unpackedLen;
				if(!unpackedLen)
					unpackedLen = file->unk20;	// what

				UInt8 * srcBuf = new UInt8[file->packedLen];

				m_src->ReadBuf(srcBuf, file->packedLen);

				UInt8 * dstBuf = new UInt8[unpackedLen];

				UInt32 bytesWritten = unpackedLen;
				int decErr = uncompress(dstBuf, &bytesWritten, srcBuf, file->packedLen);
				ASSERT(decErr == Z_OK);
				ASSERT(bytesWritten == unpackedLen);

				dst.WriteBuf(dstBuf, unpackedLen);

				delete [] dstBuf;
				delete [] srcBuf;
			}
			else
			{
				IDataStream::CopySubStreams(&dst, m_src, file->unpackedLen);
				_MESSAGE("%08X uncompressed (%s)", i, dstPath.c_str());
			}
		}
	}
}

void Archive::Extract_DX10(const char * dstRoot)
{
	for(UInt32 i = 0; i < m_textures.size(); i++)
	{
		Texture * tex = &m_textures[i];

		std::string dstPath = dstRoot;
		dstPath += "\\";
		dstPath += m_names[i];

		IFileStream::MakeAllDirs(dstPath.c_str());

		IFileStream dst;
		if(dst.Create(dstPath.c_str()))
		{
			_DMESSAGE("%08X: %08X %.4s %08X %02X %02X %04X %04X %04X %02X %02X %04X", i,
				tex->hdr.nameHash, tex->hdr.ext, tex->hdr.dirHash,
				tex->hdr.unk0C, tex->hdr.numChunks, tex->hdr.chunkHdrLen,
				tex->hdr.width, tex->hdr.height, tex->hdr.numMips, tex->hdr.format, tex->hdr.unk16);

			DDS_HEADER ddsHeader = { 0 };

			ddsHeader.dwSize = sizeof(ddsHeader);
			ddsHeader.dwHeaderFlags = DDS_HEADER_FLAGS_TEXTURE | DDS_HEADER_FLAGS_LINEARSIZE | DDS_HEADER_FLAGS_MIPMAP;
			ddsHeader.dwHeight = tex->hdr.height;
			ddsHeader.dwWidth = tex->hdr.width;
			ddsHeader.dwMipMapCount = tex->hdr.numMips;
			ddsHeader.ddspf.dwSize = sizeof(DDS_PIXELFORMAT);
			ddsHeader.dwSurfaceFlags = DDS_SURFACE_FLAGS_TEXTURE | DDS_SURFACE_FLAGS_MIPMAP;

			bool ok = true;

			switch(tex->hdr.format)
			{
			case DXGI_FORMAT_BC1_UNORM:
				ddsHeader.ddspf.dwFlags = DDS_FOURCC;
				ddsHeader.ddspf.dwFourCC = MAKEFOURCC('D', 'X', 'T', '1');
				ddsHeader.dwPitchOrLinearSize = tex->hdr.width * tex->hdr.height / 2;	// 4bpp
				break;

			case DXGI_FORMAT_BC2_UNORM:
				ddsHeader.ddspf.dwFlags = DDS_FOURCC;
				ddsHeader.ddspf.dwFourCC = MAKEFOURCC('D', 'X', 'T', '3');
				ddsHeader.dwPitchOrLinearSize = tex->hdr.width * tex->hdr.height;	// 8bpp
				break;

			case DXGI_FORMAT_BC3_UNORM:
				ddsHeader.ddspf.dwFlags = DDS_FOURCC;
				ddsHeader.ddspf.dwFourCC = MAKEFOURCC('D', 'X', 'T', '5');
				ddsHeader.dwPitchOrLinearSize = tex->hdr.width * tex->hdr.height;	// 8bpp
				break;

			case DXGI_FORMAT_BC5_UNORM:
				ddsHeader.ddspf.dwFlags = DDS_FOURCC;
				if(m_useATIFourCC)
					ddsHeader.ddspf.dwFourCC = MAKEFOURCC('A', 'T', 'I', '2');	// this is more correct but the only thing I have found that supports it is the nvidia photoshop plugin
				else
					ddsHeader.ddspf.dwFourCC = MAKEFOURCC('D', 'X', 'T', '5');

				ddsHeader.dwPitchOrLinearSize = tex->hdr.width * tex->hdr.height;	// 8bpp
				break;

			case DXGI_FORMAT_BC7_UNORM:
				// totally wrong but not worth writing out the DX10 header
				ddsHeader.ddspf.dwFlags = DDS_FOURCC;
				ddsHeader.ddspf.dwFourCC = MAKEFOURCC('B', 'C', '7', '\0');
				ddsHeader.dwPitchOrLinearSize = tex->hdr.width * tex->hdr.height;	// 8bpp
				break;

			case DXGI_FORMAT_B8G8R8A8_UNORM:
				ddsHeader.ddspf.dwFlags = DDS_RGBA;
				ddsHeader.ddspf.dwRGBBitCount = 32;
				ddsHeader.ddspf.dwRBitMask =	0x00FF0000;
				ddsHeader.ddspf.dwGBitMask =	0x0000FF00;
				ddsHeader.ddspf.dwBBitMask =	0x000000FF;
				ddsHeader.ddspf.dwABitMask =	0xFF000000;
				ddsHeader.dwPitchOrLinearSize = tex->hdr.width * tex->hdr.height * 4;	// 32bpp
				break;

			case DXGI_FORMAT_R8_UNORM:
				ddsHeader.ddspf.dwFlags = DDS_RGB;
				ddsHeader.ddspf.dwRGBBitCount = 8;
				ddsHeader.ddspf.dwRBitMask =	0xFF;
				ddsHeader.dwPitchOrLinearSize = tex->hdr.width * tex->hdr.height;	// 8bpp
				break;

			default:
				_ERROR("unhandled format %02X (%d) (%s)", tex->hdr.format, tex->hdr.format, dstPath.c_str());
				ok = false;
				break;
			}

			if(ok)
			{
				dst.Write32(DDS_MAGIC);	// 'DDS '
				dst.WriteBuf(&ddsHeader, sizeof(ddsHeader));

				gLog.Indent();

				for(UInt32 j = 0; j < tex->chunks.size(); j++)
				{
					DX10Chunk * chunk = &tex->chunks[j];

					_DMESSAGE("%016I64X %08X %08X %04X %04X %08X",
						chunk->offset, chunk->packedLen, chunk->unpackedLen,
						chunk->startMip, chunk->endMip, chunk->unk14);

					UInt8 * srcBuf = new UInt8[chunk->packedLen];

					m_src->SetOffset(chunk->offset);
					m_src->ReadBuf(srcBuf, chunk->packedLen);

					UInt8 * dstBuf = new UInt8[chunk->unpackedLen];

					UInt32 bytesWritten = chunk->unpackedLen;
					int decErr = uncompress(dstBuf, &bytesWritten, srcBuf, chunk->packedLen);
					ASSERT(decErr == Z_OK);
					ASSERT(bytesWritten == chunk->unpackedLen);

					dst.WriteBuf(dstBuf, chunk->unpackedLen);

					delete [] dstBuf;
					delete [] srcBuf;
				}

				gLog.Outdent();
			}
		}
	}
}