//////////////////////////////////////////////////////////////////////////
//
//  Tga image translucency fixer
//  Created by Wernaeh
//
//////////////////////////////////////////////////////////////////////////


// Windows stuff
#include "Windows.h"

// Files
#include <fstream>


// For shortening stuff, directly include
// the std namespace
using namespace std;


// The global exit flag for the message loop
bool bExit = false;

// The main programm window
HWND hMainWin = NULL;


// The name of the current targa file
char *pFileName = NULL;

// The current targa file data
char *pPixels = NULL;
int iResX = 0;
int iResY = 0;
char cBypp = 0;


// The default (supported) targa header
const char cTgaHeader[12] =
	{ 0, 0, 2, 0,  0, 0, 0, 0,  0, 0, 0, 0 };


// Present a message box with the specified text
// indicating that some problem occured.
void MsgBox(const char *text)
{
	MessageBox(hMainWin, text, "Problem:",
			   MB_OK | MB_ICONEXCLAMATION);
}

	
// Load a targa file with the given file name
// This method returns false if the targa file can not
// be modified, but the error is ignorable. This includes
// invalid file types, or tga's with corrupt headers.
// If everything is ok with the file, the global image data
// vars are filled.
bool LoadTga()
{
	ifstream file;

	// Open the targa file
	file.open(pFileName, ios::in | ios::binary | ios::_Nocreate);

	if (!file)
	{
		MsgBox("Failed to open file for reading");
		return false;
	}

	// Extract the file header
	char header[12];

	if (!file.read(header, sizeof(char) * 12))
	{
		MsgBox("Corrupt tga header");
		return false;
	}

	// Only support uncompressed targas
	if (memcmp(header, cTgaHeader, 12) != 0)
	{
		MsgBox("Unsupported targa header");
		return false;
	}

	// Extract file resolution
	short resx;
	short resy;

	if (!file.read((char*)&resx, sizeof(short)) ||
		!file.read((char*)&resy, sizeof(short)))
	{
		MsgBox("Corrupt targa resolution");
		return false;
	}

	// Check for valid resolution
	if (resx <= 0 || resx > 2048 ||
		resy <= 0 || resy > 2048)
	{
		MsgBox("Targa resolution not supported");
		return false;
	}

	// Extract bitdepth
	char bitdepth;

	if (!file.read(&bitdepth, sizeof(char)))
	{
		MsgBox("Corrupt targa bitdepth");
		return false;
	}

	// Only open targas with 24 or 32 bits per pixel
	if (bitdepth != 32 && bitdepth != 24)
	{
		MsgBox("Targa bitdepth not supported");
		return false;
	}

	// Jump over last part of header
	file.seekg(sizeof(char), ios::cur);

	// Apply image sizing and depth information
	iResX = resx;
	iResY = resy;	
	cBypp = bitdepth / 8;

	// Reserve image memory
	pPixels =
		(char*)malloc(iResX * iResY * cBypp * sizeof(char));

	if (!pPixels)
	{
		throw("Failed to allocate pixel memory");
	}

	// Load actual image from file
	if (!file.read(pPixels, iResX * iResY * cBypp * sizeof(char)))
	{
		MsgBox("Unexpected end in image");
		return false;
	}

	// Close file again
	file.close();

	return true;
}


// Process the targa image in memory
void ProcessImage()
{
	// No need to process targa files without
	// alpha channel
	if (cBypp != 4)
	{
		return;
	}

	// Create process map
	char *map =
		(char*)malloc(iResX * iResY * sizeof(char));

	if (!map)
	{
		throw("Failed to alloc temporary memory");
	}

	// Initialize process map
	int pixel = 3;

	for (int i = 0;
		 i < iResX * iResY;
		 ++i)
	{
		if (pPixels[pixel] != 0)
		{
			map[i] = 0;
		}
		else
		{
			map[i] = 2;
		}

		pixel += 4;
	}
		
	// Correct pixels until each one has been
	// treated.
	bool anyfound;

	do
	{
		anyfound = false;

		// Set all pixels neighbouring to at least
		// a single treated pixel to the average of
		// all treated neighbours.
		int index = 0;

		for (int i = 0;
			 i < iResY;
			 ++i)
		{
			for (int j = 0;
				 j < iResX;
				 ++j)
			{

				// Only look at untreated pixels
				if (map[index] == 2)
				{

					// Build coords for neighbor pixels
					int xnext = j + 1;
					if (xnext >= iResX)
					{
						xnext -= iResX;
					}

					int xprev = j - 1;
					if (xprev < 0)
					{
						xprev += iResX;
					}

					// Y coords for neighbors
					int ynext = i + 1;
					if (ynext >= iResY)
					{
						ynext -= iResY;
					}

					int yprev = i - 1;
					if (yprev < 0)
					{
						yprev += iResY;
					}

					// Build neighbour indices
					int indices[8] =
						{ ynext * iResX + xnext,
						  ynext * iResX + j,
						  ynext * iResX + xprev,
						  i * iResX + xnext,
						  i * iResX + xprev,
						  yprev * iResX + xnext,
						  yprev * iResX + j,
						  yprev * iResX + xprev };

					// Temporary accumulators
					int accumr = 0;
					int accumg = 0;
					int accumb = 0;
					int accnt = 0;

					// Accumulate all completed neighbors
					for (int k = 0;
						 k < 8;
						 ++k)
					{
						if (map[indices[k]] == 0)
						{
							accumr += pPixels[indices[k] * 4];
							accumg += pPixels[indices[k] * 4 + 1];
							accumb += pPixels[indices[k] * 4 + 2];
							++accnt;
						}
					}

					// Assign result of accumulation and update
					// the process map
					if (accnt != 0)
					{
						pPixels[index * 4] = accumr / accnt;
						pPixels[index * 4 + 1] = accumg / accnt;
						pPixels[index * 4 + 2] = accumb / accnt;

						map[index] = 1;
					}
				}

				++index;
			}
		}

		// Mark all treated pixels as completed
		// Also raise anyfound flag if any pixel
		// was treated.
		for (int i = 0;
			 i < iResX * iResY;
			 ++i)
		{
			if (map[i] == 1)
			{
				map[i] = 0;
				anyfound = true;
			}
		}
	}
	while (anyfound);

	// Release process map
	free(map);
}


// Write the targa file back to disk
void SaveTga()
{
	ofstream file;

	// Open the file for output
	file.open(pFileName, ios::out | ios::binary | ios::trunc); 

	if (!file)
	{
		MsgBox("Failed to open file for saving");
		return;
	}

	// Create some temporaries
	char bitdepth = cBypp * 8;
	char bituse = 0;

	// Write out all targa image data
	file.write(cTgaHeader, 12 * sizeof(char));

	file.write((char*)&iResX, sizeof(short));
	file.write((char*)&iResY, sizeof(short));
	file.write(&bitdepth, sizeof(char));
	file.write(&bituse, sizeof(char));

	file.write(pPixels, iResX * iResY * cBypp * sizeof(char));

	// Close file again
	file.close();
}


// Release all image data
void ReleaseImage()
{
	iResX = 0;
	iResY = 0;
	cBypp = 0;

	if (pPixels)
	{
		free(pPixels);
		pPixels = NULL;
	}
}


// Process a single file dropped into the main window
void ProcessFile()
{
	// Load the targa file
	if (LoadTga())
	{
		// Fix the image data
		ProcessImage();

		// Write it back into the file
		SaveTga();
	}

	// Release all image data again
	ReleaseImage();
}


// Process a number of files dropped into the window
void ProcessDrop(HDROP drop)
{
	// Retrieve file count dropped
	int fcount =
		DragQueryFile(drop, 0xffffffff, NULL, 0);

	for (int i = 0;
		 i < fcount;
		 ++i)
	{
		// Allocate buffer for file name
		int namlen =
			DragQueryFile(drop, i, NULL, 0) + 1;

		pFileName =
			(char*)malloc(sizeof(char) * namlen);

		if (!pFileName)
		{
			throw("Failed to allocate memory");
		}

		// Retrieve file name
		DragQueryFile(drop, i, pFileName, namlen);
				
		// Proces the specified file
		ProcessFile();

		// Release name buffer
		free(pFileName);
	}

	// Clean up memory occupied by the drop
	DragFinish(drop);

	// Bring up a message box to indicate
	// that we finished
	MessageBox(hMainWin, "Files fixed",
			   "Finished", MB_OK);
}


// The window procedure
LRESULT CALLBACK WndProc
	(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
	// Handle incoming message
	switch (uMsg)
	{
		case WM_CLOSE:
			// Force program to terminate
			bExit = true;
		return 0;  

		case WM_DROPFILES:
			// Some files were dropped, so process them
			ProcessDrop((HDROP)wParam);
		return 0;
	}

	// Let default window procedure handle all
	// unhandled messages.
	return
		DefWindowProc(hWnd, uMsg, wParam, lParam);
}


// Run the actual program
void Run()
{
	// Repeat main loop until exit flag is raised
	while (!bExit)
	{

		// Handle all incoming windows messages
		MSG msg;

		while (PeekMessage(&msg, hMainWin, 0, 0, PM_REMOVE))
		{
			TranslateMessage(&msg);
			DispatchMessage(&msg);
		}
	}
}


// The programm creation method
void Create()
{
	// Create program window class
	WNDCLASS wndcls =
		{ 0, &WndProc, 0, 0, GetModuleHandle(NULL), NULL,
		  LoadCursor(NULL, IDC_ARROW), GetSysColorBrush(COLOR_3DFACE),
		  NULL, "TgaFixWindow" };

	if (RegisterClass(&wndcls) == 0)
	{
		throw("Failed to register window class");
	}

	// Create program window
	hMainWin = CreateWindowEx
		(WS_EX_APPWINDOW | WS_EX_ACCEPTFILES | WS_EX_TOPMOST |
		 WS_EX_CLIENTEDGE, "TgaFixWindow", "TgaFix",
		 WS_CAPTION | WS_BORDER | WS_POPUP | WS_SYSMENU,
		 50, 50, 150, 100, NULL, NULL, GetModuleHandle(NULL), NULL);

	if (hMainWin == NULL)
	{
		throw("Failed to create window");
	}

	// Create the static "drag items here" text
	HWND tempwin = CreateWindow
		("STATIC", "Drag TGA files here", WS_CHILD | WS_VISIBLE,
		 22, 20, 100, 100, hMainWin, NULL, GetModuleHandle(NULL), NULL);

	if (tempwin == NULL)
	{
		throw("Failed to create static window text");
	}

	// Make window visible
	ShowWindow(hMainWin, SW_SHOW);
}


// Programm shutdown method
void Destroy()
{
	// Release image data in the case of an exceptional
	// shutdown.
	ReleaseImage();

	// Destroy programm window
	if (hMainWin)
	{
		DestroyWindow(hMainWin);
	}

	// Unregister window class
	UnregisterClass("TgaFixWindow", GetModuleHandle(NULL));
}


// This programm is intended to fix the problems occuring with
// alpha channel images on certain imaging applications. Photoshop
// for example sets all completely translucent pixels to white, which
// later on gives problems with mipmapping. This tga fixer allows the
// user to specifiy a number of files which then are all repaired.
// Repair is performed in that all fully transparent pixels are recolored
// with the average color of their (non translucent) neighbours.
// This process continues iteratively until all translucent pixels have
// been recolored appropriately. Note since this is only a tool, the
// code ain't as clean as it should perhaps be.
int CALLBACK WinMain
	(IN HINSTANCE hInstance, IN HINSTANCE hPrevInstance,
	 IN LPSTR lpCmdLine, IN int nShowCmd)
{
	try
	{
		// Setup out program on startup
		Create();

		try
		{
			// Run program
			Run();
		}
		catch (const char *except)
		{
			MsgBox(except);
		}

		// Clean up after we finished
		Destroy();
	}
	catch (const char *except)
	{
		MsgBox(except);
	}
}