/*
 * WBTitle
 *
 * Version 1.0
 *
 * Public Domain Software
 *
 * This program replaces the Amiga's Workbench title bar so that it
 * shows all types of memory available: Chip, Public, VM, and Retina.
 */

#include <string.h>
#include <ctype.h>

#include <exec/types.h>
#include <exec/memory.h>
#include <intuition/intuition.h>
#include <dos/dos.h>

#include <proto/exec.h>
#include <proto/intuition.h>
#include <proto/dos.h>
#include <proto/retina.h>
#include <proto/vmm.h>

#include <clib/alib_protos.h>

// The SetWindowTitles function offset
#define SWTOffset -276

enum {
	JMPINSTR = 0x4ef9
};

typedef struct JmpEntry {
	UWORD Instr;
	APTR Func;
} JmpEntry;

static BOOL Replace(void);
static void Restore(void);
static void __asm new_SetWindowTitles(register __a0 struct Window *,
								register __a1 UBYTE *, register __a2 UBYTE *,
								register __a6 struct Library *);
static void NumberToString(ULONG number);
static void format_number(ULONG number);
static void no_lead_triple(ULONG number);
static void lead_triple(ULONG number);

// Global variables
struct RetinaBase *RetinaBase;
struct Library *VMMBase;

// Local variables
static char verstring[] = "$VER: WBTitle 1.0 " __AMIGADATE__;
static char port_name[] = "WBTitle";

static void __asm (*old_SetWindowTitles)(register __a0 struct Window *,
									register __a1 UBYTE *,
									register __a2 UBYTE *,
									register __a6 struct Library *);
static JmpEntry *SWTEntry;
static char *chip_str, *fast_str, *public_str, *total_str, *virtual_str;
static char *retina_str, *memorder_str, *prefix_str, *suffix_str, *labels_str;
static char *comma_str, *mem_str, *last_str, *buff_ptr;
static char num_buf[14], comma;
static int units, num_size, mem_str_len, prefix_str_len;
static BOOL sep1000, labels_after;

/*
 * The main method for replacing an Amiga OS function as safe as
 * possible is to place the function with a jump table that is
 * allocated.  While the function is replaced, the jump table simply
 * jumps to my routine:
 *
 * jmp  _new_SetWindowTitles
 *
 * When the user asks the program to quit, we can't simply put the
 * pointer back that SetFunction() gives us since someone else might
 * have replaced the function.  So, we first see if the pointer we
 * get back points to the out jump table.  If so, then we _can_ put
 * the pointer back like normal (no one has replaced the function
 * while we has it replaced).  But if the pointer isn't mine, then
 * we have to replace the jump table function pointer to the old
 * function pointer:
 *
 * jmp  _old_SetWindowTitles
 *
 * Finally, we only deallocate the jump table _if_ we did not have
 * to change the jump table.
 */

main(int argc, char *argv[])
{
	struct MsgPort *port;

	// FindPort() Forbid()
	Forbid();

	port = FindPort(port_name);
	if (port) {
		struct MsgPort *reply_port;

		// Create a reply port
		reply_port = CreateMsgPort();
		if (reply_port) {
			struct Message msg;

			// Set fields in message structure
			msg.mn_ReplyPort = reply_port;
			msg.mn_Length = sizeof(struct Message);

			// Send the message
			PutMsg(port, &msg);

			// Finished with port, so stop FindPort() Forbid()
			Permit();

			// Wait for a reply
			do {
				WaitPort(reply_port);
			} while (GetMsg(reply_port) == NULL);

			// Clear and Delete reply_port Forbid()
			Forbid();

			// Clear any messages
			while (GetMsg(reply_port));

			// Delete the reply port
			DeleteMsgPort(reply_port);

			// Clear and Delete reply_port stop Forbid()
			Permit();
		} else {
			// Finished with port, so stop FindPort() Forbid()
			Permit();
		}
	} else if (port = CreateMsgPort()) {
		struct Message *msg;
		char **ttypes;

		// Finished with port, so stop FindPort() Forbid()
		Permit();

		// Setup quitting port
		port->mp_Node.ln_Name = port_name;
		port->mp_Node.ln_Pri = -120;

		// Add quitting port to public list
		AddPort(port);

		// Open the Retina library
		RetinaBase = (struct RetinaBase *)OpenLibrary("retina.library", 0);

		// Check on VMM; look for port, and alloc signal
		if (FindPort("VMM_Port")) {
			VMMBase = OpenLibrary("vmm.library", 0);
		}

		// Setup to read some arguements
		ttypes = ArgArrayInit(argc, argv);

		// Read some arguments
		prefix_str = ArgString(ttypes, "PREFIX", "Amiga Workbench  ");
		suffix_str = ArgString(ttypes, "SUFFIX", "");
		labels_str = ArgString(ttypes, "LABELS", "AFTER");

		chip_str = ArgString(ttypes, "CHIP", " Chip  ");
		fast_str = ArgString(ttypes, "FAST", " Fast  ");
		public_str = ArgString(ttypes, "PUBLIC", " Public  ");
		total_str = ArgString(ttypes, "TOTAL", " Total  ");
		virtual_str = ArgString(ttypes, "VIRTUAL", " Virtual  ");
		retina_str = ArgString(ttypes, "RETINA", " Retina  ");

		memorder_str = ArgString(ttypes, "MEMORDER", "CVPR");
		sep1000 = (ArgString(ttypes, "THOUSANDSEP", NULL) != NULL);
		units = ArgInt(ttypes, "UNITS", 1);
		comma_str = ArgString(ttypes, "SEPERATOR", NULL);
		if (comma_str && comma_str[0]) {
			comma = comma_str[0];
		} else {
			comma = ',';
		}

		labels_after = (stricmp(labels_str, "BEFORE") != 0);
		if (units < 0) {
			units = 1;
		}

		// Compute the size of numbers
		{
			ULONG maxnum = 0xffffffff;

			// Compute the number of digits
			maxnum /= units;
			while (maxnum) {
				maxnum /= 10;
				num_size++;
			}

			// Compute the number of commas
			if (sep1000) {
				int com = num_size;
				while (com > 3) {
					num_size++;
					com -= 3;
				}
			}
		}

		// Compute the size of the mem string
		{
			ULONG i = 0;
			char ch;

			// Add in the prefix string size
			mem_str_len = prefix_str_len = strlen(prefix_str);

			// Add in the memory parts
			while (ch = toupper(memorder_str[i])) {
				switch (ch) {
					case 'C':
						mem_str_len += strlen(chip_str) + num_size;
						break;
					case 'F':
						mem_str_len += strlen(fast_str) + num_size;
						break;
					case 'P':
						mem_str_len += strlen(public_str) + num_size;
						break;
					case 'T':
						mem_str_len += strlen(total_str) + num_size;
						break;
					case 'V':
						if (VMMBase) {
							mem_str_len += strlen(virtual_str) + num_size;
						}
						break;
					case 'R':
						if (RetinaBase) {
							mem_str_len += strlen(retina_str) + num_size;
						}
						break;
				}
				i++;
			}

			// Add in the suffix string size
			mem_str_len += strlen(suffix_str);
		}

		if (mem_str_len > 0) {
			// Allocate the memory
			mem_str = AllocMem(mem_str_len + 1, MEMF_PUBLIC);
			if (mem_str) {
				strcpy(mem_str, prefix_str);
				last_str = &mem_str[prefix_str_len];

				// Attempt to replace function
				if (Replace()) {
					// Wait for someone to signal me to quit
					do {
						WaitPort(port);
						msg = GetMsg(port);
					} while (msg == NULL);
					ReplyMsg(msg);

					// Restore function
					Restore();
				}
				FreeMem(mem_str, mem_str_len + 1);
			}
		}

		// Cleanup from reading arguments
		ArgArrayDone();

		// Remove port from public access
		RemPort(port);

		// Clear and Delete port Forbid()
		Forbid();

		// Clear the port of messages
		while (msg = GetMsg(port)) {
			ReplyMsg(msg);
		}

		// Closedown quitting port
		DeleteMsgPort(port);

		// Clear and Delete port stop Forbid()
		Permit();
	}
}

static BOOL Replace(void)
{
	// Allocate the jump table
	SWTEntry = AllocMem(sizeof(JmpEntry), 0);
	if (SWTEntry) {
		// Replacement Forbid()
		Forbid();

		// Replace the function with pointer to jump table
		old_SetWindowTitles = SetFunction((struct Library *)IntuitionBase,
											SWTOffset, (ULONG (*)())SWTEntry);

		// Setup the jump table
		SWTEntry->Instr = JMPINSTR;
		SWTEntry->Func = new_SetWindowTitles;

		// Clear the cpu's cache so the execution cache is valid
		CacheClearU();

		// Stop the replacement Forbid()
		Permit();

		return TRUE;
	} else {
		return FALSE;
	}
}

static void Restore(void)
{
	BOOL my_table;
	ULONG (*func)();

	// Fix back Forbid()
	Forbid();

	// Put old pointer back and get current pointer at same time
	func = SetFunction((struct Library *)IntuitionBase, SWTOffset,
							(ULONG (*)())old_SetWindowTitles);

	// Check to see if the pointer we get back is ours
	if ((JmpEntry *)func != SWTEntry) {
		// If not, leave jump table in place
		my_table = FALSE;
		SetFunction((struct Library *)IntuitionBase, SWTOffset,
					func);
		SWTEntry->Func = old_SetWindowTitles;
	} else {
		// If so, free the jump table
		my_table = TRUE;
		FreeMem(SWTEntry, sizeof(JmpEntry));
	}

	// Clear the cpu's cache so the execution cache is valid
	CacheClearU();

	// Stop fix back Forbid()
	Permit();

	// Let the user know if the jump table couldn't be freed
	if (!my_table) {
		DisplayBeep(NULL);
	}

	// Wait 5 seconds to try and guarantee that all tasks have
	// finished executing inside my replacement function before
	// quitting.  There's no real way to guarantee, though.
	Delay(250);
}

static void __saveds __asm new_SetWindowTitles(
							register __a0 struct Window *window,
							register __a1 UBYTE *win_title,
							register __a2 UBYTE *scr_title,
							register __a6 struct Library *lib)
{
	if (scr_title && (scr_title != (UBYTE *)0xffffffff) &&
			(*scr_title == 'A') && strstr(scr_title, "graphics mem")) {
		ULONG i = 0;
		char ch;

		// Set last_str to nothing
		*last_str = 0;

		// Add in the memory parts
		while (ch = toupper(memorder_str[i])) {
			switch (ch) {
				case 'C':
					if (!labels_after) {
						strcat(last_str, chip_str);
					}
					NumberToString(AvailMem(MEMF_CHIP));
					strcat(last_str, num_buf);
					if (labels_after) {
						strcat(last_str, chip_str);
					}
					break;
				case 'F':
					if (!labels_after) {
						strcat(last_str, fast_str);
					}
					NumberToString(AvailMem(MEMF_FAST));
					strcat(last_str, num_buf);
					if (labels_after) {
						strcat(last_str, fast_str);
					}
					break;
				case 'P':
					if (!labels_after) {
						strcat(last_str, public_str);
					}
					NumberToString(AvailMem(MEMF_PUBLIC));
					strcat(last_str, num_buf);
					if (labels_after) {
						strcat(last_str, public_str);
					}
					break;
				case 'T':
					if (!labels_after) {
						strcat(last_str, total_str);
					}
					NumberToString(AvailMem(0));
					strcat(last_str, num_buf);
					if (labels_after) {
						strcat(last_str, total_str);
					}
					break;
				case 'V':
					if (VMMBase) {
						if (!labels_after) {
							strcat(last_str, virtual_str);
						}
						NumberToString(AvailVMem(0));
						strcat(last_str, num_buf);
						if (labels_after) {
							strcat(last_str, virtual_str);
						}
					}
					break;
				case 'R':
					if (RetinaBase) {
						if (!labels_after) {
							strcat(last_str, retina_str);
						}
						NumberToString(Retina_AvailMem(0));
						strcat(last_str, num_buf);
						if (labels_after) {
							strcat(last_str, retina_str);
						}
					}
					break;
			}
			i++;
		}
		strcat(last_str, suffix_str);
		old_SetWindowTitles(window, win_title, mem_str, lib);
	} else {
		old_SetWindowTitles(window, win_title, scr_title, lib);
	}
}

/*
 * Convert the number to a string in the format the user specified.
 * The string generated is 0 terminated.
 *
 * Uses global variables: buff_ptr, num_buff
 */

static void NumberToString(ULONG number)
{
	number /= units;
	if (sep1000) {
		buff_ptr = num_buf;
		format_number(number);
		*buff_ptr = 0;
	} else {
		stcul_d(num_buf, number);
	}
}

static void format_number(ULONG number)
{
	if (number > 1000) {
		format_number(number / 1000);
		lead_triple(number % 1000);
	} else {
		no_lead_triple(number);
	}
}

static void no_lead_triple(ULONG number)
{
	if (number > 10) {
		no_lead_triple(number / 10);
		*buff_ptr++ = ((number % 10) | 0x30);
	} else {
		*buff_ptr++ = (number | 0x30);
	}
}

static void lead_triple(ULONG number)
{
	*buff_ptr++ = comma;
	*buff_ptr++ = ((number / 100) | 0x30);
	number %= 100;
	*buff_ptr++ = ((number / 10) | 0x30);
	number %= 10;
	*buff_ptr++ = (number | 0x30);
}