#include "ImGui.hpp"

#include <algorithm>
#include <cassert>
#include <cctype>
#include <cmath>
#include <format>
#include <string>
#include <string_view>
#include <vector>

#include <raylib.h>

namespace {

struct CodepointSpan {
	u32 codepoint {};
	usize start {};
	usize end {};
};

auto decode_utf8(std::string_view text) -> std::vector<CodepointSpan>
{
	std::vector<CodepointSpan> spans;
	usize i = 0;
	spans.reserve(text.size());

	while (i < text.size()) {
		u8 const byte = static_cast<u8>(text[i]);
		usize const start = i;
		usize length = 1;
		u32 cp = 0xFFFD;

		if (byte < 0x80) {
			cp = byte;
		} else if ((byte & 0xE0) == 0xC0) {
			if (i + 1 < text.size()) {
				u8 const b1 = static_cast<u8>(text[i + 1]);
				if ((b1 & 0xC0) == 0x80) {
					u32 const t = ((static_cast<u32>(byte) & 0x1F) << 6)
					    | (static_cast<u32>(b1) & 0x3F);
					if (t >= 0x80) {
						cp = t;
						length = 2;
					}
				}
			}
		} else if ((byte & 0xF0) == 0xE0) {
			if (i + 2 < text.size()) {
				u8 const b1 = static_cast<u8>(text[i + 1]);
				u8 const b2 = static_cast<u8>(text[i + 2]);
				if ((b1 & 0xC0) == 0x80 && (b2 & 0xC0) == 0x80) {
					u32 const t = ((static_cast<u32>(byte) & 0x0F) << 12)
					    | ((static_cast<u32>(b1) & 0x3F) << 6)
					    | (static_cast<u32>(b2) & 0x3F);
					if (t >= 0x800 && (t < 0xD800 || t > 0xDFFF)) {
						cp = t;
						length = 3;
					}
				}
			}
		} else if ((byte & 0xF8) == 0xF0) {
			if (i + 3 < text.size()) {
				u8 const b1 = static_cast<u8>(text[i + 1]);
				u8 const b2 = static_cast<u8>(text[i + 2]);
				u8 const b3 = static_cast<u8>(text[i + 3]);
				if ((b1 & 0xC0) == 0x80 && (b2 & 0xC0) == 0x80
				    && (b3 & 0xC0) == 0x80) {
					u32 const t = ((static_cast<u32>(byte) & 0x07) << 18)
					    | ((static_cast<u32>(b1) & 0x3F) << 12)
					    | ((static_cast<u32>(b2) & 0x3F) << 6)
					    | (static_cast<u32>(b3) & 0x3F);
					if (t >= 0x10000 && t <= 0x10FFFF) {
						cp = t;
						length = 4;
					}
				}
			}
		}

		spans.push_back(CodepointSpan { cp, start, start + length });
		i += length;
	}

	return spans;
}

auto encode_utf8(u32 cp) -> std::string
{
	char buf[5] = { 0, 0, 0, 0, 0 };
	int len = 0;
	if (cp <= 0x7F) {
		buf[len++] = static_cast<char>(cp);
	} else if (cp <= 0x7FF) {
		buf[len++] = static_cast<char>(0xC0 | (cp >> 6));
		buf[len++] = static_cast<char>(0x80 | (cp & 0x3F));
	} else if (cp <= 0xFFFF) {
		if (cp >= 0xD800 && cp <= 0xDFFF)
			return {};
		buf[len++] = static_cast<char>(0xE0 | (cp >> 12));
		buf[len++] = static_cast<char>(0x80 | ((cp >> 6) & 0x3F));
		buf[len++] = static_cast<char>(0x80 | (cp & 0x3F));
	} else if (cp <= 0x10FFFF) {
		buf[len++] = static_cast<char>(0xF0 | (cp >> 18));
		buf[len++] = static_cast<char>(0x80 | ((cp >> 12) & 0x3F));
		buf[len++] = static_cast<char>(0x80 | ((cp >> 6) & 0x3F));
		buf[len++] = static_cast<char>(0x80 | (cp & 0x3F));
	} else {
		return {};
	}
	return std::string(buf, len);
}

auto rune_index_for_byte(std::string_view text, usize byte_offset) -> int
{
	auto spans { decode_utf8(text) };
	int idx = 0;
	for (auto const &span : spans) {
		if (span.start >= byte_offset)
			break;
		idx++;
	}
	if (byte_offset >= text.size())
		idx = static_cast<int>(spans.size());
	return idx;
}

auto clamp_preedit_index(int value, usize text_size) -> usize
{
	if (value < 0)
		return 0;
	auto const as_size { static_cast<usize>(value) };
	return std::min(as_size, text_size);
}

auto slice_bytes(std::string_view text, usize begin, usize end)
    -> std::string_view
{
	if (begin > text.size())
		begin = text.size();
	if (end > text.size())
		end = text.size();
	if (end < begin)
		end = begin;
	return std::string_view(text.data() + begin, end - begin);
}

constexpr float HORIZONTAL_PADDING = 6.0f;
constexpr float VERTICAL_PADDING = 4.0f;
constexpr float CARET_WIDTH = 2.0f;
constexpr double CARET_BLINK_INTERVAL = 0.5;
constexpr float CARET_DESCENT_FRACTION = 0.25f;

} // namespace

ImGui::ImGui(std::shared_ptr<TextRenderer> text_renderer)
    : m_text_renderer(text_renderer)
{
}

void ImGui::begin(u32 const rune, bool ctrl, bool shift)
{
	m_rune = rune;
	m_ctrl = ctrl;
	m_shift = shift;
}

void ImGui::end() { }

void ImGui::set_font(FontHandle font) { m_font = font; }

auto ImGui::focused_text_input() const -> std::optional<usize>
{
	if (m_focused_id == 0)
		return std::nullopt;
	return m_focused_id;
}

auto ImGui::text_input_surrounding(usize id, std::pmr::string const &str) const
    -> std::optional<TextInputSurrounding>
{
	auto it { m_ti_states.find(id) };
	if (it == m_ti_states.end())
		return std::nullopt;
	TextInputSurrounding info;
	info.text.assign(str.data(), str.size());
	info.caret_byte = std::min(it->second.caret_byte, str.size());
	info.cursor = static_cast<int>(info.caret_byte);
	info.anchor = static_cast<int>(info.caret_byte);
	return info;
}

auto ImGui::text_input_cursor(usize id) const -> std::optional<TextInputCursor>
{
	auto it { m_ti_states.find(id) };
	if (it == m_ti_states.end())
		return std::nullopt;
	TextInputCursor cursor;
	cursor.rect = it->second.caret_rect;
	cursor.visible
	    = it->second.caret_visible && !it->second.preedit_cursor_hidden;
	return cursor;
}

void ImGui::ime_commit_text(std::pmr::string &str, std::string_view text)
{
	if (m_focused_id == 0)
		return;
	auto it { m_ti_states.find(m_focused_id) };
	if (it == m_ti_states.end())
		return;
	auto &state { it->second };
	usize insert_pos = std::min(state.caret_byte, str.size());
	if (!text.empty())
		str.insert(insert_pos, text);
	state.caret_byte = insert_pos + text.size();
	std::string_view const view(str.data(), str.size());
	state.current_rune_idx = rune_index_for_byte(view, state.caret_byte);
	state.caret_timer = 0.0;
	state.caret_visible = true;
	state.external_change = true;
}

void ImGui::ime_delete_surrounding(
    std::pmr::string &str, usize before, usize after)
{
	if (m_focused_id == 0)
		return;
	auto it { m_ti_states.find(m_focused_id) };
	if (it == m_ti_states.end())
		return;
	auto &state { it->second };
	usize caret_byte = std::min(state.caret_byte, str.size());
	usize start = before > caret_byte ? 0 : caret_byte - before;
	usize end = std::min(caret_byte + after, str.size());
	if (end > start) {
		str.erase(start, end - start);
		state.caret_byte = start;
		std::string_view const view(str.data(), str.size());
		state.current_rune_idx = rune_index_for_byte(view, state.caret_byte);
		state.caret_timer = 0.0;
		state.caret_visible = true;
		state.external_change = true;
	}
}

void ImGui::ime_set_preedit(std::string text, int cursor_begin, int cursor_end)
{
	if (m_focused_id == 0)
		return;
	auto it { m_ti_states.find(m_focused_id) };
	if (it == m_ti_states.end())
		return;
	auto &state { it->second };
	state.preedit_text = std::move(text);
	state.preedit_cursor_hidden = (cursor_begin == -1 && cursor_end == -1);
	usize const size = state.preedit_text.size();
	if (state.preedit_cursor_hidden) {
		state.preedit_cursor_begin = 0;
		state.preedit_cursor_end = 0;
	} else {
		auto begin_clamped { clamp_preedit_index(cursor_begin, size) };
		auto end_clamped { clamp_preedit_index(cursor_end, size) };
		state.preedit_cursor_begin = static_cast<int>(begin_clamped);
		state.preedit_cursor_end = static_cast<int>(end_clamped);
	}
	state.preedit_active
	    = !state.preedit_text.empty() || !state.preedit_cursor_hidden;
	if (state.preedit_active) {
		state.caret_timer = 0.0;
		state.caret_visible = !state.preedit_cursor_hidden;
	}
}

void ImGui::ime_clear_preedit()
{
	if (m_focused_id == 0)
		return;
	auto it { m_ti_states.find(m_focused_id) };
	if (it == m_ti_states.end())
		return;
	auto &state { it->second };
	state.preedit_text.clear();
	state.preedit_cursor_begin = 0;
	state.preedit_cursor_end = 0;
	state.preedit_active = false;
	state.preedit_cursor_hidden = false;
	state.caret_visible = true;
	state.caret_timer = 0.0;
}

auto ImGui::text_input(usize id, std::pmr::string &str, Rectangle rec,
    TextInputOptions options) -> std::bitset<2>
{
	assert(id != 0);
	assert(
	    m_font.has_value() && "ImGui font must be set before using text input");

	bool submitted { false };
	bool changed { false };

	auto &state { m_ti_states[id] };

	assert(!options.multiline && "Multiline not yet implemented.");

	if (m_focused_id == 0)
		m_focused_id = id;

	if (style().font_size > rec.height) {
		TraceLog(LOG_WARNING,
		    std::format("Text size for text input {} is bigger than height ({} "
		                "> {}). Clipping will occur.",
		        id, style().font_size, rec.height)
		        .c_str());
	}

	std::string_view str_view(str.data(), str.size());
	auto spans { decode_utf8(str_view) };

	auto is_space { [](u32 cp) -> bool {
		if (cp == '\n' || cp == '\r' || cp == '\t' || cp == '\v' || cp == '\f')
			return true;
		if (cp <= 0x7F)
			return std::isspace(static_cast<unsigned char>(cp)) != 0;
		return false;
	} };

	auto clamp_cursor { [&]() -> usize {
		int const max_idx = static_cast<int>(spans.size());
		state.current_rune_idx = std::clamp(state.current_rune_idx, 0, max_idx);
		if (state.current_rune_idx == max_idx)
			return str.size();
		return spans[state.current_rune_idx].start;
	} };

	usize caret_byte = clamp_cursor();

	auto refresh_spans { [&]() {
		str_view = std::string_view(str.data(), str.size());
		spans = decode_utf8(str_view);
		caret_byte = clamp_cursor();
	} };

	auto erase_range { [&](usize byte_begin, usize byte_end) {
		if (byte_end > byte_begin && byte_begin < str.size()) {
			str.erase(byte_begin, byte_end - byte_begin);
			changed = true;
		}
	} };

	bool caret_activity = false;

	if (m_focused_id == id && m_rune != 0) {
		bool request_refresh = false;
		auto handle_backspace { [&]() {
			if (state.current_rune_idx <= 0
			    || state.current_rune_idx > static_cast<int>(spans.size()))
				return;

			if (m_ctrl) {
				int idx = state.current_rune_idx;
				int scan = idx - 1;
				while (scan >= 0
				    && is_space(spans[static_cast<usize>(scan)].codepoint))
					scan--;
				while (scan >= 0
				    && !is_space(spans[static_cast<usize>(scan)].codepoint))
					scan--;
				int start_idx = std::max(scan + 1, 0);
				usize byte_begin = spans[static_cast<usize>(start_idx)].start;
				usize byte_end = (idx >= static_cast<int>(spans.size()))
				    ? str.size()
				    : spans[static_cast<usize>(idx)].start;
				erase_range(byte_begin, byte_end);
				state.current_rune_idx = start_idx;
			} else {
				auto const &prev
				    = spans[static_cast<usize>(state.current_rune_idx - 1)];
				erase_range(prev.start, prev.end);
				state.current_rune_idx--;
			}
			request_refresh = true;
		} };

		auto handle_delete { [&]() {
			if (state.current_rune_idx < 0
			    || state.current_rune_idx >= static_cast<int>(spans.size())) {
				if (!m_ctrl)
					return;
			}

			int idx = state.current_rune_idx;
			if (m_ctrl) {
				int scan = idx;
				while (scan < static_cast<int>(spans.size())
				    && is_space(spans[static_cast<usize>(scan)].codepoint))
					scan++;
				while (scan < static_cast<int>(spans.size())
				    && !is_space(spans[static_cast<usize>(scan)].codepoint))
					scan++;
				usize byte_begin = (idx < static_cast<int>(spans.size()))
				    ? spans[static_cast<usize>(idx)].start
				    : str.size();
				usize byte_end = (scan < static_cast<int>(spans.size()))
				    ? spans[static_cast<usize>(scan)].start
				    : str.size();
				erase_range(byte_begin, byte_end);
			} else if (idx < static_cast<int>(spans.size())) {
				auto const &curr { spans[static_cast<usize>(idx)] };
				erase_range(curr.start, curr.end);
			}
			request_refresh = true;
		} };

		switch (m_rune) {
		case 1: // Left (H)
			if (state.current_rune_idx > 0) {
				state.current_rune_idx--;
				if (m_ctrl) {
					while (state.current_rune_idx > 0
					    && is_space(
					        spans[static_cast<usize>(state.current_rune_idx)]
					            .codepoint))
						state.current_rune_idx--;
					while (state.current_rune_idx > 0
					    && !is_space(
					        spans[static_cast<usize>(state.current_rune_idx)]
					            .codepoint))
						state.current_rune_idx--;
				}
				caret_byte = clamp_cursor();
			}
			break;
		case 4: // Right (L)
			if (state.current_rune_idx < static_cast<int>(spans.size())) {
				state.current_rune_idx++;
				if (m_ctrl) {
					while (
					    state.current_rune_idx < static_cast<int>(spans.size())
					    && is_space(spans[static_cast<usize>(
					                          state.current_rune_idx - 1)]
					            .codepoint))
						state.current_rune_idx++;
					while (
					    state.current_rune_idx < static_cast<int>(spans.size())
					    && !is_space(spans[static_cast<usize>(
					                           state.current_rune_idx - 1)]
					            .codepoint))
						state.current_rune_idx++;
				}
				caret_byte = clamp_cursor();
			}
			break;
		case 3: // Up (K)
			state.current_rune_idx = 0;
			caret_byte = clamp_cursor();
			break;
		case 2: // Down (J)
			state.current_rune_idx = static_cast<int>(spans.size());
			caret_byte = clamp_cursor();
			break;
		case 8: // Backspace
			handle_backspace();
			break;
		case 0x7F: // Delete
			handle_delete();
			break;
		case '\r':
		case '\n':
			if (options.multiline) {
				auto encoded { encode_utf8('\n') };
				if (!encoded.empty()) {
					str.insert(caret_byte, encoded);
					state.current_rune_idx++;
					changed = true;
					request_refresh = true;
				}
			} else {
				submitted = true;
			}
			break;
		default:
			if (m_rune >= 0x20) {
				auto encoded { encode_utf8(m_rune) };
				if (!encoded.empty()) {
					str.insert(caret_byte, encoded);
					state.current_rune_idx++;
					changed = true;
					request_refresh = true;
				}
			}
			break;
		}

		if (request_refresh) {
			refresh_spans();
		} else {
			caret_byte = clamp_cursor();
		}
		caret_activity = true;
	}

	state.caret_byte = caret_byte;

	double const dt = static_cast<double>(GetFrameTime());
	if (m_focused_id == id) {
		if (state.preedit_active && state.preedit_cursor_hidden) {
			state.caret_visible = false;
			state.caret_timer = 0.0;
		} else if (caret_activity) {
			state.caret_timer = 0.0;
			state.caret_visible = true;
		} else {
			if (state.caret_timer == 0.0)
				state.caret_visible = true;
			state.caret_timer += dt;
			if (state.caret_timer >= CARET_BLINK_INTERVAL) {
				int toggles = static_cast<int>(
				    state.caret_timer / CARET_BLINK_INTERVAL);
				state.caret_timer
				    = std::fmod(state.caret_timer, CARET_BLINK_INTERVAL);
				if (toggles % 2 == 1)
					state.caret_visible = !state.caret_visible;
			}
		}
	} else {
		state.caret_visible = false;
		state.caret_timer = 0.0;
	}

	Vector2 prefix_metrics { 0.0f, 0.0f };
	Vector2 full_metrics { 0.0f, 0.0f };
	Vector2 preedit_metrics { 0.0f, 0.0f };
	Vector2 caret_preedit_metrics { 0.0f, 0.0f };
	Vector2 selection_prefix_metrics { 0.0f, 0.0f };
	Vector2 selection_metrics { 0.0f, 0.0f };
	std::string display_buffer;
	bool const has_preedit = state.preedit_active
	    && (!state.preedit_text.empty() || !state.preedit_cursor_hidden);

	if (m_font.has_value() && m_text_renderer) {
		int const font_px = static_cast<int>(style().font_size);
		std::string_view const prefix_view(str.data(), caret_byte);
		prefix_metrics
		    = m_text_renderer->measure_text(*m_font, prefix_view, font_px);

		if (has_preedit) {
			std::string_view const preedit_view(
			    state.preedit_text.data(), state.preedit_text.size());
			preedit_metrics
			    = m_text_renderer->measure_text(*m_font, preedit_view, font_px);

			auto caret_idx { clamp_preedit_index(
				state.preedit_cursor_end, preedit_view.size()) };
			caret_preedit_metrics = m_text_renderer->measure_text(
			    *m_font, slice_bytes(preedit_view, 0, caret_idx), font_px)

			                            if (!state.preedit_cursor_hidden
			                                && state.preedit_cursor_begin
			                                    != state.preedit_cursor_end)
			{
				auto sel_begin
				    = clamp_preedit_index(std::min(state.preedit_cursor_begin,
				                              state.preedit_cursor_end),
				        preedit_view.size());
				auto sel_end
				    = clamp_preedit_index(std::max(state.preedit_cursor_begin,
				                              state.preedit_cursor_end),
				        preedit_view.size());
				selection_prefix_metrics = m_text_renderer->measure_text(
				    *m_font, slice_bytes(preedit_view, 0, sel_begin), font_px);
				selection_metrics = m_text_renderer->measure_text(*m_font,
				    slice_bytes(preedit_view, sel_begin, sel_end), font_px);
			}

			display_buffer.reserve(str.size() + state.preedit_text.size());
			display_buffer.append(prefix_view);
			display_buffer.append(preedit_view);
			display_buffer.append(str_view.substr(caret_byte));
			full_metrics = m_text_renderer->measure_text(*m_font,
			    std::string_view(display_buffer.data(), display_buffer.size()),
			    font_px);
		} else {
			full_metrics
			    = m_text_renderer->measure_text(*m_font, str_view, font_px);
		}
	}

	float caret_offset = prefix_metrics.x + caret_preedit_metrics.x;
	state.cursor_position.x = caret_offset;

	float const available_width
	    = std::max(0.0f, rec.width - 2.0f * HORIZONTAL_PADDING);
	if (full_metrics.x <= available_width) {
		state.scroll_offset.x = 0.0f;
	} else {
		float &scroll = state.scroll_offset.x;
		float caret_local = caret_offset - scroll;
		if (caret_local > available_width) {
			scroll = caret_offset - available_width;
		} else if (caret_local < 0.0f) {
			scroll = caret_offset;
		}
		scroll = std::clamp(
		    scroll, 0.0f, std::max(0.0f, full_metrics.x - available_width));
	}
	state.scroll_offset.y = 0.0f;

	Color const bg_col { 16, 16, 16, 100 };
	Color const border_col { 220, 220, 220, 180 };
	DrawRectangleRec(rec, bg_col);
	DrawRectangleLinesEx(rec, 1.0f, border_col);

	float const text_top = rec.y + VERTICAL_PADDING;
	float const baseline_y = text_top + style().font_size;
	float const max_caret_height
	    = std::max(0.0f, rec.height - 2.0f * VERTICAL_PADDING);
	float caret_height = style().font_size * (1.0f + CARET_DESCENT_FRACTION);
	caret_height = std::min(caret_height,
	    max_caret_height > 0.0f ? max_caret_height : caret_height);
	if (caret_height <= 0.0f)
		caret_height = style().font_size;
	float caret_top = baseline_y - style().font_size;
	float const min_top = rec.y + VERTICAL_PADDING;
	float const max_top = rec.y + rec.height - VERTICAL_PADDING - caret_height;
	caret_top = std::clamp(caret_top, min_top, max_top);
	float caret_bottom = caret_top + caret_height;
	float const desired_bottom
	    = baseline_y + style().font_size * CARET_DESCENT_FRACTION;
	if (caret_bottom < desired_bottom) {
		float const adjust = desired_bottom - caret_bottom;
		caret_top = std::min(caret_top + adjust, max_top);
		caret_bottom = caret_top + caret_height;
	}
	state.cursor_position.y = caret_top;

	float const caret_draw_x
	    = rec.x + HORIZONTAL_PADDING + caret_offset - state.scroll_offset.x;
	state.caret_rect = Rectangle {
		caret_draw_x,
		caret_top,
		CARET_WIDTH,
		caret_height,
	};

	BeginScissorMode(rec.x, rec.y, rec.width, rec.height);
	{
		if (m_font.has_value() && m_text_renderer) {
			int const font_px = static_cast<int>(style().font_size);
			Vector2 const base_pos {
				rec.x + HORIZONTAL_PADDING - state.scroll_offset.x,
				baseline_y,
			};
			Color const &text_color = style().text_color;
			std::string_view const left_view(str.data(), caret_byte);
			std::string_view const right_view(
			    str.data() + caret_byte, str.size() - caret_byte);

			m_text_renderer->draw_text(
			    *m_font, left_view, base_pos, font_px, text_color);

			float advance = prefix_metrics.x;
			if (has_preedit) {
				Vector2 const preedit_pos {
					base_pos.x + advance,
					base_pos.y,
				};
				if (selection_metrics.x > 0.0f) {
					float const sel_offset
					    = prefix_metrics.x + selection_prefix_metrics.x;
					Rectangle const sel_rect {
						rec.x + HORIZONTAL_PADDING + sel_offset
						    - state.scroll_offset.x,
						caret_top,
						selection_metrics.x,
						caret_height,
					};
					Color const highlight { style().selection_color };
					DrawRectangleRec(sel_rect, highlight);
				}
				Color const &preedit_color { selection_metrics.x > 0.0f
					    ? style().selection_text_color
					    : style().preedit_color };
				m_text_renderer->draw_text(*m_font,
				    std::string_view(
				        state.preedit_text.data(), state.preedit_text.size()),
				    preedit_pos, font_px, preedit_color);
				advance += preedit_metrics.x;
			}

			Vector2 const right_pos {
				base_pos.x + advance,
				base_pos.y,
			};
			m_text_renderer->draw_text(
			    *m_font, right_view, right_pos, font_px, text_color);

			if (m_focused_id == id && state.caret_visible) {
				Rectangle caret_rect = state.caret_rect;
				caret_rect.x = std::round(caret_rect.x);
				Color const caret_color { text_color };
				DrawRectangleRec(caret_rect, caret_color);
			}
		}
	}
	EndScissorMode();

	if (state.external_change) {
		changed = true;
		state.external_change = false;
	}

	return std::bitset<2> { static_cast<unsigned long long>(
		(submitted ? 1 : 0) | (changed ? 2 : 0)) };
}

auto ImGui::list_view(usize id, Rectangle bounds, usize elements,
    std::function<Vector2(usize i)> draw_cb, ListViewOptions options) -> bool
{
	auto &state { m_lv_states[id] };

	bool submitted { false };

	bool select_next = m_next_lv_next;
	m_next_lv_next = false;
	bool select_previous = m_next_lv_previous;
	m_next_lv_previous = false;
	bool select_clear = m_next_lv_clear;
	m_next_lv_clear = false;

	BeginScissorMode(bounds.x, bounds.y, bounds.width, bounds.height);

	EndScissorMode();

	m_prev_lv_selected_item = state.selected_item;

	return submitted;
}