#include "tray_icon.h"

#define TRAY_WINDOW_MESSAGE (WM_USER + 101)

namespace {
	// A map that never holds allocated memory when it is empty. This map will be created with placement new as a static variable,
	// and its destructor will be never called, and it shouldn't leak memory if it contains no items at program exit.
	// This dirty trick is useful when you create your trayicon object as static. In this case we can not control the
	// order of destruction of this map object and the static trayicon object. However this dirty trick ensures that
	// the map is freed exactly when the destructor of the last static trayicon is unregistering itself.
	class TrayIconMap {
	public:
		typedef UINT KeyType;
		typedef TrayIcon *ValueType;

		//typedef StdMap std::map<KeyType, ValueType>;
		struct StdMap : public std::map<KeyType, ValueType>{};
		typedef StdMap::iterator iterator;

		TrayIconMap() : m_empty(true) {}

		ValueType &operator[](KeyType k) {
			return getOrCreateStdMap()[k];
		}

		ValueType *find(KeyType k) {
			if (m_empty) {
				return false;
			}
			StdMap::iterator it = getStdMap().find(k);
			if (it == getStdMap().end()) {
				return NULL;
			}		
			return &it->second;
		}

		int erase(KeyType k) {
			if (m_empty) {
				return 0;
			}		
			StdMap &m = getStdMap();
			int res = (int)m.erase(k);
			if (m.empty()) {
				m.~StdMap();
				m_empty = true;
			}
			return res;
		}

		bool empty() const {
			return m_empty;
		}

		iterator begin() {
			assert(!m_empty);
			return m_empty ? iterator() : getStdMap().begin();
		}

		iterator end() {
			assert(!m_empty);
			return m_empty ? iterator() : getStdMap().end();
		}

	private:
		StdMap &getStdMap() {
			assert(!m_empty);
			return (StdMap &)m_mapBuffer;
		}

		StdMap &getOrCreateStdMap() {
			if (m_empty) {
				new ((void *)&m_mapBuffer) StdMap();
				m_empty = false;
			}
			return (StdMap &)m_mapBuffer;
		}

	private:
		bool m_empty;
		char m_mapBuffer[sizeof(StdMap)];
	};

	static TrayIconMap &getTrayIconMap() {
		// This hack prevents running the destructor of our map, so it isn't problem if someone tries to reach this from a static destructor.
		// Because of using MyMap this will not cause a memory leak if the user removes all items from the container before exiting.
		static char tray_icon_buffer[sizeof(TrayIconMap)];
		static bool initialized = false;
		if (!initialized) {
			initialized = true;
			new ((void *)tray_icon_buffer) TrayIconMap();
		}
		return (TrayIconMap &)tray_icon_buffer;
	}

	static UINT getNextTrayIconId() {
		static UINT next_id = 1;
		return next_id++;
	}

}


static const UINT g_WndMsgTaskbarCreated = RegisterWindowMessage(TEXT("TaskbarCreated"));
LRESULT CALLBACK TrayIcon::MessageProcessorWndProc(HWND hWnd, UINT uMsg, WPARAM wParam, LPARAM lParam)
{
	if (uMsg == TRAY_WINDOW_MESSAGE) {
		if (TrayIcon **ppIcon = getTrayIconMap().find((UINT)wParam)) {
			(*ppIcon)->onMessage((UINT)lParam);
		}
		return 0;
	} else if (uMsg == g_WndMsgTaskbarCreated) {
		TrayIconMap &m = getTrayIconMap();
		if (!m.empty()) {
			for (std::map<UINT, TrayIcon *>::const_iterator it = m.begin(), eit = m.end(); it != eit; ++it) {
				TrayIcon *pTrayIcon = it->second;
				pTrayIcon->onTaskbarCreated();
			}
		}
	}
	return DefWindowProc(hWnd, uMsg, wParam, lParam);
}

HWND TrayIcon::GetMessageProcessorHWND()
{
	static HWND hWnd = NULL;
	if (!hWnd) {
		static const TCHAR TRAY_ICON_MESSAGE_PROCESSOR_WND_CLASSNAME[] = TEXT("TRAY_ICON_MESSAGE_PROCESSOR_WND_CLASS");
		HINSTANCE hInstance = (HINSTANCE)GetModuleHandle(NULL);

		WNDCLASSEX wc;
		wc.cbSize = sizeof(wc);
		wc.cbClsExtra = 0;
		wc.cbWndExtra = 0;
		wc.hbrBackground = (HBRUSH)(COLOR_WINDOW + 1);
		wc.hCursor = LoadCursor(NULL, IDC_ARROW);
		wc.hIcon = LoadIcon(NULL, IDI_WINLOGO);
		wc.hIconSm = NULL;
		wc.hInstance = hInstance;
		wc.lpfnWndProc = MessageProcessorWndProc;
		wc.lpszClassName = TRAY_ICON_MESSAGE_PROCESSOR_WND_CLASSNAME;
		wc.lpszMenuName = NULL;
		wc.style = 0;
		if (!RegisterClassEx(&wc)) {
			return NULL;
		}

		hWnd = CreateWindowEx(
			0,
			TRAY_ICON_MESSAGE_PROCESSOR_WND_CLASSNAME,
			TEXT("TRAY_ICON_MESSAGE_PROCESSOR_WND"),
			WS_POPUP,
			0, 0, 0, 0,
			NULL,
			NULL,
			hInstance,
			NULL);
	}
	return hWnd;
}

TrayIcon::TrayIcon(LPCWSTR name, bool visible, HICON hIcon, bool destroyIconInDestructor)
	: m_id(getNextTrayIconId())
	, m_name(name)
	, m_hIcon(hIcon)
	, m_visible(false)
	, m_destroyIconInDestructor(destroyIconInDestructor)
	, m_messageFunc(NULL)
	, m_listener(NULL)
	, m_userData(NULL)
{
	getTrayIconMap()[m_id] = this;
	setVisible(visible);
}

TrayIcon::~TrayIcon()
{
	setVisible(false);
	setIcon(NULL, m_destroyIconInDestructor);
	getTrayIconMap().erase(m_id);
}

HICON TrayIcon::internalGetIcon() const
{
	return m_hIcon ? m_hIcon : ::LoadIcon(NULL, IDI_APPLICATION);
}

bool TrayIcon::addIcon()
{
	NOTIFYICONDATAW data;
	fillNotifyIconData(data);
	data.uFlags |= NIF_MESSAGE | NIF_ICON | NIF_TIP;
	data.uCallbackMessage = TRAY_WINDOW_MESSAGE;
	data.hIcon = internalGetIcon();

	size_t tip_len = std::max(sizeof(data.szTip) - 1, wcslen(m_name.c_str()));
	wmemcpy(data.szTip, m_name.c_str(), tip_len);
	data.szTip[tip_len] = 0;

	return FALSE != Shell_NotifyIcon(NIM_ADD, &data);
}

bool TrayIcon::removeIcon()
{
	NOTIFYICONDATAW data;
	fillNotifyIconData(data);
	return FALSE != Shell_NotifyIcon(NIM_DELETE, &data);
}

void TrayIcon::onTaskbarCreated()
{
	if (m_visible) {
		addIcon();
	}
}

void TrayIcon::setName(LPCWSTR name)
{
	m_name = name;
	if (m_visible) {
		NOTIFYICONDATAW data;
		fillNotifyIconData(data);
		data.uFlags |= NIF_TIP;

		size_t tip_len = std::max(sizeof(data.szTip) - 1, wcslen(name));
		wmemcpy(data.szTip, name, tip_len);
		data.szTip[tip_len] = 0;

		Shell_NotifyIcon(NIM_MODIFY, &data);
	}
}

bool TrayIcon::setVisible(bool visible)
{
	if (m_visible == visible) {
		return true;
	}

	m_visible = visible;
	if (m_visible) {
		return addIcon();
	}
	return removeIcon();
}

void TrayIcon::setIcon(HICON hNewIcon, bool destroyCurrentIcon)
{
	if (m_hIcon == hNewIcon) {
		return;
	}
	if (destroyCurrentIcon && m_hIcon) {
		DestroyIcon(m_hIcon);
	}
	m_hIcon = hNewIcon;

	if (m_visible) {
		NOTIFYICONDATAW data;
		fillNotifyIconData(data);
		data.uFlags |= NIF_ICON;
		data.hIcon = internalGetIcon();
		Shell_NotifyIcon(NIM_MODIFY, &data);
	}
}

bool TrayIcon::notify(LPCWSTR title, LPCWSTR msg, NotificationIcon icon)
{
#ifndef NOTIFYICONDATA_V2_SIZE
	return false;
#else
	if (!m_visible) {
		return false;
	}

	NOTIFYICONDATAW data;
	fillNotifyIconData(data);
	data.cbSize = NOTIFYICONDATAW_V2_SIZE;
	data.uFlags |= NIF_INFO;
	data.dwInfoFlags = icon;
	data.uTimeout = 10000;

	wcscpy_s(data.szInfoTitle, title);
	wcscpy_s(data.szInfo, msg);

	return FALSE != Shell_NotifyIcon(NIM_MODIFY, &data);
#endif
}

void TrayIcon::onMessage(UINT uMsg)
{
	if (m_messageFunc) {
		m_messageFunc(this, uMsg);
	}

	if (m_listener) {
		m_listener->onTrayIconMessage(this, uMsg);
	}
}

void TrayIcon::fillNotifyIconData(NOTIFYICONDATAW &data)
{
	memset(&data, 0, sizeof(data));
	// the basic functions need only V1
#ifdef NOTIFYICONDATA_V1_SIZE
	data.cbSize = NOTIFYICONDATA_V1_SIZE;
#else
	data.cbSize = sizeof(data);
#endif
	data.hWnd = GetMessageProcessorHWND();
	assert(data.hWnd);
	data.uID = m_id;
}