#include "frontend.h"
#include "backend/lldb/lldb_backend.h"

#include "imgui.h"
#include "imgui_internal.h"

using namespace dbgui;
using namespace dbgui::frontend;

Target::Target(std::string filename)
{
	state          = TargetState::startup;
	this->filename = filename;
	id             = 0;
	backend = std::make_shared<backend::LLDBBackend>(this->filename.c_str());
}

Frontend::Frontend()
{
	_open_popup_name_buf.resize(512);
	const char *str = "/home/klee/projects/dbgui/tmp/main";
	std::memcpy(_open_popup_name_buf.data(), str, strlen(str) + 1);
	_windows.push_back(Window::create_regs(window_id++));
	_windows.push_back(Window::create_threads(window_id++));
	_windows.push_back(Window::create_frames(window_id++));
	_windows.push_back(Window::create_disas(window_id++));
	_windows.push_back(Window::create_bp(window_id++));
	_windows.push_back(Window::create_source(window_id++));
}

void Frontend::run_frame()
{
	if (_draw_metric_window)
	{
		ImGui::ShowMetricsWindow();
	}

	if (_draw_stack_tool)
	{
		ImGui::ShowStackToolWindow();
	}

	this->handle_msgs();
	this->draw_open_popup();
	this->draw_header();

	// main window
	float height = ImGui::GetFrameHeight();

	const auto vp_size  = ImGui::GetMainViewport()->Size;
	const auto win_pos  = ImVec2{0, height * 2};
	const auto win_size = ImVec2{vp_size.x, vp_size.y - height * 3};

	const auto win_flags =
	  ImGuiWindowFlags_NoDecoration | ImGuiWindowFlags_NoNav
	  | ImGuiWindowFlags_NoMove | ImGuiWindowFlags_NoScrollWithMouse
	  | ImGuiWindowFlags_NoBringToFrontOnFocus | ImGuiWindowFlags_NoSavedSettings;

	ImGui::SetNextWindowPos(win_pos, ImGuiCond_Always);
	ImGui::SetNextWindowSize(win_size, ImGuiCond_Always);

	if (ImGui::Begin("Test", nullptr, win_flags))
	{
		// TODO: if we just do SetNextWindowDockID in each of the windows
		// this seems to crash once the first window is docked to it?
		// so some id change is necessary
		// so we should try to figure out what the dock id for the "topleft"
		// window is or whatever and use that for new windows
		// spawning them floating looks bad (and might reuse saved positions)
		this->dock_id = ImGui::GetID("MyDockSpace");
		ImGui::DockSpace(this->dock_id, ImVec2{0, 0},
		                 ImGuiDockNodeFlags_PassthruCentralNode);

		size_t window_idx = 0;
		while (window_idx < _windows.size())
		{
			auto &window = _windows[window_idx];
			if (window.draw(*this))
			{
				_windows.erase(_windows.begin() + window_idx);
			} else
			{
				window_idx++;
			}
			//this->dock_id = ImGui::GetWindowDockID();
		}
	}
	ImGui::End();

	this->draw_status();
}

void Frontend::draw_header()
{
	ImGuiWindowFlags window_flags =
	  ImGuiWindowFlags_NoScrollbar | ImGuiWindowFlags_NoSavedSettings
	  | ImGuiWindowFlags_NoDocking | ImGuiWindowFlags_MenuBar;
	float height = ImGui::GetFrameHeight();

	if (ImGui::BeginMainMenuBar())
	{
		if (ImGui::BeginMenu("File"))
		{
			if (ImGui::MenuItem("Open"))
			{
				_draw_open_popup = true;
				ImGui::OpenPopup(_open_popup_id);
			}

			ImGui::EndMenu();
		}

		if (ImGui::BeginMenu("Window"))
		{
			if (ImGui::MenuItem("Registers"))
			{
				_windows.push_back(Window::create_regs(this->window_id++));
			}
			ImGui::EndMenu();
		}

		if (ImGui::BeginMenu("Debug"))
		{
			ImGui::MenuItem("Metrics", nullptr, &_draw_metric_window);
			ImGui::MenuItem("Stack Tool", nullptr, &_draw_stack_tool);
			ImGui::EndMenu();
		}
		ImGui::EndMainMenuBar();
	}

	if (ImGui::BeginViewportSideBar("##SecondaryMenuBar", NULL, ImGuiDir_Up,
	                                height, window_flags))
	{
		if (ImGui::BeginMenuBar())
		{
			if (this->target)
			{
				switch (this->target->state)
				{
					using enum TargetState;

					case startup:
						ImGui::BeginDisabled();
						// buttons shouldn't flicker that way
						ImGui::Button("Continue");
						ImGui::EndDisabled();
						break;

					case stopped: ImGui::Button("Run"); break;
					case paused:
						if (ImGui::Button("Continue"))
						{
							this->target->backend->cont();
							// TODO: mark target already as running to prevent input?
						}
						break;
					case running:
						if (ImGui::Button("Pause   "))
						{
							this->target->backend->pause();
						}
						break;
				}
			} else
			{
				ImGui::BeginDisabled();
				ImGui::Button("Continue");
				ImGui::EndDisabled();
			}

			ImGui::SeparatorEx(ImGuiSeparatorFlags_Vertical, 2.f);

			const auto is_paused =
			  this->target && this->target->state == TargetState::paused;
			const auto is_running =
			  this->target && this->target->state == TargetState::running;

			if (!is_paused)
			{
				ImGui::BeginDisabled();
			}

			if (ImGui::Button("Step Over"))
			{
				if (this->target->backend->step_over(!this->target->step_instruction))
				{
					// TODO: already disable the UI into running mode
				}
			}
			if (ImGui::Button("Step Into"))
			{
				if (this->target->backend->step_into(!this->target->step_instruction))
				{
					// TODO: already disable the UI into running mode
				}
			}
			if (ImGui::Button("Step Out"))
			{
				if (this->target->backend->step_out())
				{
					// TODO: already disable the UI into running mode
				}
			}

			if (!is_paused)
			{
				ImGui::EndDisabled();
			}

			ImGui::SeparatorEx(ImGuiSeparatorFlags_Vertical, 2.f);

			if (!is_paused && !is_running)
			{
				ImGui::BeginDisabled();
			}

			ImGui::Button("Stop");
			ImGui::Button("Restart");

			if (!is_paused && !is_running)
			{
				ImGui::EndDisabled();
			}

			ImGui::SeparatorEx(ImGuiSeparatorFlags_Vertical, 2.f);

			if (!this->target)
			{
				ImGui::Button("Step Mode: Source");
			} else
			{
				if (this->target->step_instruction)
				{
					if (ImGui::Button("Step Mode: Instruction"))
					{
						this->target->step_instruction = false;
					}
				} else
				{
					if (ImGui::Button("Step Mode: Source"))
					{
						this->target->step_instruction = true;
					}
				}
			}

			ImGui::EndMenuBar();
		}
		ImGui::End();
	}
}

void Frontend::draw_status()
{
	ImGuiWindowFlags window_flags =
	  ImGuiWindowFlags_NoScrollbar | ImGuiWindowFlags_NoSavedSettings
	  | ImGuiWindowFlags_NoDocking | ImGuiWindowFlags_MenuBar;
	float height = ImGui::GetFrameHeight();

	if (ImGui::BeginViewportSideBar("##MainStatusBar", NULL, ImGuiDir_Down,
	                                height, window_flags))
	{
		if (ImGui::BeginMenuBar())
		{
			ImGui::Text("Happy status bar");
			ImGui::EndMenuBar();
		}
	}
	ImGui::End();
}

void Frontend::draw_open_popup()
{
	if (!_open_popup_id)
	{
		_open_popup_id = ImGui::GetID("Open Executable##OpenPopup");
	}
	if (ImGui::BeginPopupModal("Open Executable##OpenPopup", &_draw_open_popup,
	                           ImGuiWindowFlags_AlwaysAutoResize))
	{
		ImGui::InputText("File Path", _open_popup_name_buf.data(),
		                 _open_popup_name_buf.size());

		if (ImGui::Button("Start"))
		{
			this->target = Target{_open_popup_name_buf.data()};
			this->target->backend->start();
			ImGui::CloseCurrentPopup();
		}
		ImGui::EndPopup();
	}
}

void Frontend::handle_msgs()
{
	if (!this->target)
	{
		return;
	}

	while (true)
	{
		auto opt = this->target->backend->retrieve_msg_for_frontend();
		if (!opt)
		{
			break;
		}
		auto *msg = opt->get();

		printf("Got msg %u\n", msg->type);

		switch (msg->type)
		{
			using enum BackToFront::MsgType;

			case state_change:
			{
				const auto &info = std::get<BackToFront::StateChange>(msg->data);
				this->handle_state_change(info);
				break;
			}
			case ip_change:
				// TODO
				break;
			case initial_proc_info:
				this->handle_proc_info(
				  std::move(std::get<BackToFront::InitialProcessInfo>(msg->data)));
				break;
			case regs_changed:
				this->handle_reg_change(std::get<BackToFront::RegsChanged>(msg->data));
				break;
			case thread_changed:
				this->handle_thread_change(
				  std::move(std::get<BackToFront::ThreadChange>(msg->data)));
				break;
			case thread_removed:
				this->handle_thread_remove(
				  std::get<BackToFront::ThreadRemoved>(msg->data));
				break;
			case frame_changed:
				this->handle_frame_change(
				  std::move(std::get<BackToFront::FrameChanged>(msg->data)));
				break;
			case frame_removed:
				this->handle_frame_remove(
				  std::get<BackToFront::FrameRemoved>(msg->data));
				break;
			case data_result:
			{
				const auto &result = std::get<BackToFront::DataResult>(msg->data);
				for (size_t i = 0; i < result.nodes.size(); ++i)
				{
					uint16_t idx = result.nodes[i].idx;
					if (this->target->data_res_nodes.size() <= idx)
					{
						this->target->data_res_nodes.resize(idx + 1);
					}
					this->target->data_res_nodes[idx] = std::move(result.nodes[i]);
				}
				if (result.src_node_id)
				{
					auto id    = result.src_node_id->second;
					auto found = false;
					for (auto &entry : this->target->src_id_to_data_idx)
					{
						if (entry.first == id)
						{
							entry.second = result.src_node_id->first;
							found        = true;
							break;
						}
					}
					if (!found)
					{
						this->target->src_id_to_data_idx.push_back(
						  std::make_pair(id, result.src_node_id->first));
					}
					for (auto &window : _windows)
					{
						window.handle_source_updated(*this->target, id);
					}
				}
				break;
			}
			case selected_frame_changed:
			{
				if (this->target)
				{
					this->target->selected_frame =
					  std::get<BackToFront::SelectedFrameChanged>(msg->data).idx;
				}
				break;
			}
			case selected_thread_changed:
			{
				if (this->target)
				{
					this->target->selected_thread =
					  std::get<BackToFront::SelectedThreadChanged>(msg->data).idx;
				}
				break;
			}
		}
	}
}

void Frontend::handle_state_change(const BackToFront::StateChange &info)
{
	this->target->state = info.new_state;

	// TODO: display in status bar
	printf("State changed to %u because %u\n", info.new_state, info.reason);
}

void Frontend::handle_proc_info(BackToFront::InitialProcessInfo &&info)
{
	this->target->arch = info.arch;
	this->target->id   = info.pid;

	auto &sets = this->target->reg_sets;
	printf("RegSet Size: %lu\n", info.reg_sets.size());
	for (size_t set_idx = 0; set_idx < info.reg_sets.size(); ++set_idx)
	{
		sets.emplace_back();
		auto &set = sets.back();
		printf("Got set %s\n", info.reg_sets[set_idx].name.c_str());
		set.name = std::move(info.reg_sets[set_idx].name);

		for (size_t i = 0; i < info.reg_sets[set_idx].reg_names.size(); ++i)
		{
			set.regs.emplace_back();
			printf(" Got reg %s\n", info.reg_sets[set_idx].reg_names[i].c_str());
			set.regs.back().name = std::move(info.reg_sets[set_idx].reg_names[i]);
		}
	}
}

void Frontend::handle_reg_change(const BackToFront::RegsChanged &info)
{
	auto &set = this->target->reg_sets[info.set_idx];
	for (const auto &[idx, val] : info.changes)
	{
		auto &reg = set.regs[idx];
		reg.bytes.resize(val.size());

		// TODO: opt for uint64?
		std::copy(val.begin(), val.end(), reg.bytes.begin());
	}
}

void Frontend::handle_thread_remove(const BackToFront::ThreadRemoved &info)
{
	auto &threads = this->target->threads;
	if (threads.size() <= info.idx)
	{
		return;
	}

	threads[info.idx] = {};
}

void Frontend::handle_thread_change(BackToFront::ThreadChange &&info)
{
	auto &threads = this->target->threads;
	if (threads.size() <= info.idx)
	{
		threads.resize(info.idx + 1);
	}

	auto &thread = threads[info.idx];
	if (!thread)
	{
		thread = Target::Thread{
		  .id          = info.id,
		  .ip          = info.ip,
		  .stop_extra  = info.stop_extra,
		  .stop_reason = info.stop_reason,
		};
		if (info.name)
		{
			thread->name = std::move(*info.name);
		}
		if (info.cur_display_fn)
		{
			thread->cur_display_fn = std::move(*info.cur_display_fn);
		}
		return;
	}

	thread->id          = info.id;
	thread->ip          = info.ip;
	thread->stop_reason = info.stop_reason;
	thread->stop_extra  = info.stop_extra;
	if (info.name)
	{
		thread->name = std::move(*info.name);
	}
	if (info.cur_display_fn)
	{
		thread->cur_display_fn = std::move(*info.cur_display_fn);
	}
}

void Frontend::handle_frame_remove(const BackToFront::FrameRemoved &info)
{
	auto &frames = this->target->frames;
	if (frames.size() <= info.idx)
	{
		return;
	}

	frames[info.idx] = {};
}

void Frontend::handle_frame_change(BackToFront::FrameChanged &&info)
{
	auto &frames = this->target->frames;
	if (frames.size() <= info.idx)
	{
		frames.resize(info.idx + 1);
	}

	auto &frame = frames[info.idx];
	if (!frame)
	{
		printf("Adding new frame at %u with name '%s'\n", info.idx,
		       info.display_name.c_str());
		frame = Target::Frame{
		  .ip           = info.ip,
		  .display_name = info.display_name,
		};
		return;
	}

	printf("Updating frame at %u with name '%s'\n", info.idx,
	       info.display_name.c_str());
	frame->ip           = info.ip;
	frame->display_name = info.display_name;
}