#include "watchos.h"
#include "Module_UI.h"
#include "Module_Power.h"
#include <GxEPD2_BW.h>

GxEPD2_BW<GxEPD2_154_D67, GxEPD2_154_D67::HEIGHT> gfx = GxEPD2_154_D67(HW_DISPLAY_CS, HW_DISPLAY_DC, HW_DISPLAY_RESET, HW_DISPLAY_BUSY);

#define MODULE_UI_MULTICORE_STACK_SIZE 10240

struct UiWindow
{
	kernel_handle_t handle;
	IDrawable* drawable;
	byte layout_mode;
	UiWindow* parent;
	bool dirty;
	bool needs_draw;
	byte x, y, w, h;
};

struct UiFont
{
	byte id;
	GFXfont* font;
};

void destruct_window(kernel_handle_t handle, void* object)
{
}

#ifdef WATCHOS_UI_MULTICORE
TaskHandle_t draw_task;
void draw_wrapper(void* parameter)
{
	Module_UI* ui = static_cast<Module_UI*>(parameter);
	ui->do_draw();
	draw_task = nullptr;
	vTaskDelete(NULL);
}
#endif

Module_UI::Module_UI()
{
	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
		m_window[i] = nullptr;
	for (int i = 0; i < MODULE_UI_MAX_FONTS; i++)
		m_font[i] = nullptr;
}

void Module_UI::initialize()
{
	m_dirty = true;
	m_redraw = true;
}

void Module_UI::start()
{
	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
	{
		m_window[i] = nullptr;
	}

	m_root = createWindow(WATCHOS_HANDLE_NULL);
	module_power = static_cast<Module_Power*>(watchos::module(WATCHOS_MODULE_POWER));
}

void Module_UI::tick()
{
	if (m_dirty)
	{
		draw();
	}
}

kernel_handle_t Module_UI::createWindow(IDrawable* drawable, kernel_handle_t parent)
{
	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
	{
		if (m_window[i] == nullptr)
		{
			m_window[i] = new UiWindow();
			m_window[i]->drawable = drawable;
			m_window[i]->handle = watchos::kernel()->allocHandle(WATCHOS_HANDLE_TYPE_UI_WINDOW, m_window[i], (kernel_handle_destructor_t)destruct_window);
			m_window[i]->layout_mode = ui_layout_mode::NONE;
			m_window[i]->parent = nullptr;
			m_window[i]->dirty = true;
			m_window[i]->needs_draw = false;

			if (parent != WATCHOS_HANDLE_NULL)
			{
				for (int j = 0; j < MODULE_UI_MAX_WINDOWS; j++)
				{
					if (m_window[j] != nullptr && m_window[j]->handle == parent)
					{
						m_window[i]->parent = m_window[j];
						break;
					}
				}

				if (m_window[i]->parent == nullptr)
				{
					watchos::panic("Cannot find parent window handle %#010x", parent);
				}
			}

			return m_window[i]->handle;
		}
	}

	watchos::panic("Exceeded MODULE_UI_MAX_WINDOWS");
}

kernel_handle_t Module_UI::createWindow(kernel_handle_t parent)
{
	return createWindow(nullptr, parent);
}

kernel_handle_t Module_UI::getRoot()
{
	return m_root;
}

void Module_UI::setRoot(kernel_handle_t handle)
{
	m_root = handle;
	m_redraw = true;
}

void Module_UI::draw()
{
	if (m_dirty)
	{
#ifdef WATCHOS_UI_MULTICORE
		if (draw_task == nullptr)
		{
			// In multi-core mode we clear this _before_ the draw. That way if anything else is marked dirty during
			// the draw, we'll be primed to do another draw as soon as this one is complete.
			// Chances are the thing marked dirty may end up drawn during the on-going call anyway (if it hasn't been
			// drawn yet) and this will generate a extraenous draw, but that's preferable to _not_ drawing an update
			// and having out-of-date information on the screen.
			m_dirty = false;
			xTaskCreatePinnedToCore(draw_wrapper, "Draw", MODULE_UI_MULTICORE_STACK_SIZE, this, 0, &draw_task, 0);
		}
#else
		do_draw();
		// If we're doing synchronous draws then we know everything is now clean once this completes, so we mark that
		// we don't need another draw.
		m_dirty = false;
#endif
	}
}

void Module_UI::do_draw()
{
	watchos::debug("Draw!");
	module_power->disableLightSleep();
	module_power->disableDeepSleep();

	// Initialize gfx if necessary
	if (!m_initialized)
	{
		m_initialized = true;
		gfx.init(0, false);
	}

	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
	{
		if (m_window[i] != nullptr && m_window[i]->handle == m_root)
		{
			draw_window(m_window[i], 0, 0, DISPLAY_WIDTH, DISPLAY_HEIGHT);
			break;
		}
	}

	// Actually figure out what to draw
	bool whole_screen = false;
	int x, y;
	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
	{
		if (m_window[i] != nullptr && m_window[i]->needs_draw)
		{
			get_screen_coordinates(m_window[i], &x, &y);
			if (x == 0 && y == 0 && m_window[i]->w == DISPLAY_WIDTH && m_window[i]->h == DISPLAY_HEIGHT)
			{
				whole_screen = true;
				break;
			}
		}
	}

	if (whole_screen)
	{
		for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
		{
			if (m_window[i] != nullptr)
			{
				m_window[i]->needs_draw = false;
			}
		}

		gfx.display(!m_redraw);
	}
	else
	{
		for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
		{
			if (m_window[i] != nullptr && m_window[i]->needs_draw)
			{
				get_screen_coordinates(m_window[i], &x, &y);
				gfx.displayWindow(x, y, m_window[i]->w, m_window[i]->h);
				m_window[i]->needs_draw = false;
			}
		}
	}

	m_redraw = false;

	module_power->enableDeepSleep();
	module_power->enableLightSleep();
}

void Module_UI::draw_window(UiWindow* window, int x, int y, int width, int height)
{
	m_draw_x = x;
	m_draw_y = y;
	m_draw_w = width;
	m_draw_h = height;

	// Track the last calculated position for this window so we can actually figure out
	// where to update the screen.
	if (window->parent != nullptr)
	{
		window->x = x - window->parent->x;
		window->y = y - window->parent->y;
	}
	else
	{
		window->x = x;
		window->y = y;
	}
	window->w = width;
	window->h = height;

	// Draw window
	if (window->drawable != nullptr && (window->dirty || m_redraw))
	{
		window->drawable->draw(window->handle);
		window->dirty = false;
		window->needs_draw = true;
	}

	// Count children
	int child_count = 0;
	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
	{
		if (m_window[i] != nullptr && m_window[i]->parent == window)
		{
			child_count++;
		}
	}

	if (child_count == 0)
	{
		// No children to draw, we're done!
		return;
	}

	int child_offset = 0;
	int child_width, child_height;
	switch (window->layout_mode)
	{
	case ui_layout_mode::NONE:
		for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
		{
			if (m_window[i] != nullptr && m_window[i]->parent == window)
			{
				draw_window(m_window[i], x, y, width, height);
			}
		}
		break;
	case ui_layout_mode::SPLIT_HORIZONTAL:
		child_height = height / child_count;
		for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
		{
			if (m_window[i] != nullptr && m_window[i]->parent == window)
			{
				draw_window(m_window[i], x, y + (child_offset++ * child_height), width, child_height);
			}
		}
		break;
	case ui_layout_mode::SPLIT_VERTICAL:
		child_width = width / child_count;
		for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
		{
			if (m_window[i] != nullptr && m_window[i]->parent == window)
			{
				draw_window(m_window[i], x + (child_offset++ * child_width), y, child_width, height);
			}
		}
		break;
	case ui_layout_mode::ABSOLUTE:
		for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
		{
			if (m_window[i] != nullptr && m_window[i]->parent == window)
			{
				draw_window(m_window[i], x + m_window[i]->x, y + m_window[i]->y, m_window[i]->w, m_window[i]->h);
			}
		}
		break;
	}
}

void Module_UI::setLayoutMode(kernel_handle_t window_handle, ui_layout_mode layout_mode)
{
	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
	{
		if (m_window[i] != nullptr && m_window[i]->handle == window_handle)
		{
			m_window[i]->layout_mode = layout_mode;
			return;
		}
	}
	watchos::panic("Trying to set layout mode on non-existent handle: %#010x", window_handle);
}

void Module_UI::setAbsolutePosition(kernel_handle_t window_handle, byte x, byte y, byte w, byte h)
{
	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
	{
		if (m_window[i] != nullptr && m_window[i]->handle == window_handle)
		{
			m_window[i]->x = x;
			m_window[i]->y = y;
			m_window[i]->w = w;
			m_window[i]->h = h;
			return;
		}
	}

	watchos::panic("Trying to set position on non-existent handle: %#010x", window_handle);
}

void Module_UI::setDirty(kernel_handle_t window_handle)
{
	for (int i = 0; i < MODULE_UI_MAX_WINDOWS; i++)
	{
		if (m_window[i] != nullptr && m_window[i]->handle == window_handle)
		{
			m_window[i]->dirty = true;
			m_dirty = true;
			return;
		}
	}
	// TODO: Check if window is actually visible before marking m_dirty
	watchos::panic("Trying to set dirty non-existent window handle: %#010x", window_handle);
}

int Module_UI::getWidth()
{
	return m_draw_w;
}

int Module_UI::getHeight()
{
	return m_draw_h;
}

void Module_UI::fill(uint16_t colour)
{
	fill(0, 0, m_draw_w, m_draw_h, colour);
}

void Module_UI::fill(int x, int y, int width, int height, uint16_t colour)
{
	local_to_screen(&x, &y);
	gfx.fillRect(x, y, width, height, colour);
}

void Module_UI::fillRectangle(int x, int y, int w, int h, uint16_t colour)
{
	local_to_screen(&x, &y);
	gfx.fillRect(x, y, w, h, colour);
}

void Module_UI::rectangle(int x, int y, int w, int h, uint16_t colour)
{
	local_to_screen(&x, &y);
	gfx.drawRect(x, y, w, h, colour);
}

void Module_UI::line(int x0, int y0, int x1, int y1, uint16_t colour)
{
	local_to_screen(&x0, &y0);
	local_to_screen(&x1, &y1);
	gfx.drawLine(x0, y0, x1, y1, colour);
}

void Module_UI::horizontalLine(int y, uint16_t colour)
{
	int x = 0;
	local_to_screen(&x, &y);
	gfx.drawFastHLine(x, y, m_draw_w, colour);
}

void Module_UI::verticalLine(int x, uint16_t colour)
{
	int y = 0;
	local_to_screen(&x, &y);
	gfx.drawFastVLine(x, y, m_draw_h, colour);
}

void Module_UI::print(int x, int y, char* str, uint16_t colour, void* font)
{
	local_to_screen(&x, &y);
	gfx.setCursor(x, y);
	gfx.setTextColor(colour);
	gfx.setFont(static_cast<GFXfont*>(font));
	gfx.print(str);
}

void Module_UI::print(int x, int y, char* str, uint16_t colour, byte font_id)
{
	for (int i = 0; i < MODULE_UI_MAX_FONTS; i++)
	{
		if (m_font[i] != nullptr && m_font[i]->id == font_id)
		{
			print(x, y, str, colour, m_font[i]->font);
			return;
		}
	}
	watchos::panic("Unknown font id: %d", font_id);
}

void Module_UI::local_to_screen(int* x, int* y)
{
	// TODO: Check bounds?
	if (x < 0)
		*x = (m_draw_w - *x) + m_draw_x;
	else
		*x = m_draw_x + *x;

	if (y < 0)
		*y = (m_draw_h - *y) + m_draw_y;
	else
		*y = m_draw_y + *y;
}

void Module_UI::get_screen_coordinates(UiWindow* window, int* x, int* y)
{
	*x = window->x;
	*y = window->y;
	while (window->parent != nullptr)
	{
		window = window->parent;
		*x = *x + window->x;
		*y = *y + window->y;
	}
}

void Module_UI::registerFont(byte id, const void* font)
{
	for (int i = 0; i < MODULE_UI_MAX_FONTS; i++)
	{
		if (m_font[i] == nullptr)
		{
			m_font[i] = new UiFont();
			m_font[i]->id = id;
			m_font[i]->font = static_cast<GFXfont*>(const_cast<void*>(font));
			return;
		}
	}
	watchos::panic("Exceeded MODULE_UI_MAX_FONTS");
}