#include "UI.hpp"

#include "Model/GlobalStates.hpp"
#include "Model/Project.hpp"
#include "Model/Workflow/Nodes/DocumentNodes.hpp"
#include "Model/Workflow/Nodes/NumericNodes.hpp"
#include "Model/Workflow/Nodes/TextNodes.hpp"
#include "Model/Workflow/Nodes/UserInputNodes.hpp"
#include "Model/Workflow/Workflow.hpp"
#include "UI/Localization.hpp"
#include "Utils/Macros.hpp"

#include <imgui.h>
#include <imgui_node_editor.h>
#include <imgui_stdlib.h>
#include <memory>
#include <span>
#include <vector>

namespace ImNodes = ax::NodeEditor;

namespace {
enum class WorkflowCategory
{
	Numeric,
	Text,
	Documents,
	UserInput,
	SystemInput,
	Output,
};

class WorkflowDatabase
{
public:
	using WorkflowNodeConstructor = std::unique_ptr<WorkflowNode> (*)();

	struct Candidate
	{
		WorkflowNodeConstructor Constructor;
		std::string Name;
		WorkflowCategory Category;
	};

private:
	std::vector<Candidate> mCandidates;

	int mTextOffset;
	int mDocumentOffset;
	int mUserInputOffset;
	int mSystemInputNodes;
	int mOutputOffset;

public:
	// clang-format off
	std::span<Candidate> GetNumericNodes() { return { &mCandidates[0], (size_t)(mTextOffset - 0) }; };
	std::span<Candidate> GetTextNodes() { return { &mCandidates[mTextOffset], (size_t)(mDocumentOffset - mTextOffset) }; };
	std::span<Candidate> GetDocumentNodes() { return { &mCandidates[mDocumentOffset], (size_t)(mUserInputOffset - mDocumentOffset) }; };
	std::span<Candidate> GetUserInputNodes() { return { &mCandidates[mUserInputOffset], (size_t)(mSystemInputNodes - mUserInputOffset) }; };
	std::span<Candidate> GetSystemInputNodes() { return { &mCandidates[mSystemInputNodes], (size_t)(mOutputOffset - mSystemInputNodes) }; };
	std::span<Candidate> GetOutputNodes() { return { &mCandidates[mOutputOffset], (size_t)(mCandidates.size() - mOutputOffset) }; };
	// clang-format on

public:
	static WorkflowDatabase& GetInstance()
	{
		static WorkflowDatabase database;
		return database;
	}

public:
	WorkflowDatabase()
	{
		// Numeric nodes offset start at 0
		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<NumericOperationNode>(NumericOperationNode::Addition); },
			.Name = "Add",
			.Category = WorkflowCategory::Numeric,
		});
		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<NumericOperationNode>(NumericOperationNode::Subtraction); },
			.Name = "Subtract",
			.Category = WorkflowCategory::Numeric,
		});
		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<NumericOperationNode>(NumericOperationNode::Multiplication); },
			.Name = "Multiply",
			.Category = WorkflowCategory::Numeric,
		});
		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<NumericOperationNode>(NumericOperationNode::Division); },
			.Name = "Divide",
			.Category = WorkflowCategory::Numeric,
		});
		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<NumericExpressionNode>(); },
			.Name = "Evaluate expression",
			.Category = WorkflowCategory::Numeric,
		});

		mTextOffset = mCandidates.size();
		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<TextFormatterNode>(); },
			.Name = "Fill template text",
			.Category = WorkflowCategory::Text,
		});

		mDocumentOffset = mCandidates.size();
		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<DocumentTemplateExpansionNode>(); },
			.Name = "Document template",
			.Category = WorkflowCategory::Documents,
		});

		/* Inputs */

		mUserInputOffset = mCandidates.size();
		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<FormInputNode>(); },
			.Name = "Input: form",
			.Category = WorkflowCategory::UserInput,
		});

		mCandidates.push_back(Candidate{
			.Constructor = []() -> std::unique_ptr<WorkflowNode> { return std::make_unique<DatabaseRowsInputNode>(); },
			.Name = "Input: database rows",
			.Category = WorkflowCategory::UserInput,
		});

		mSystemInputNodes = mCandidates.size();

		/* Outputs */

		mOutputOffset = mCandidates.size();
	}
};

class WorkflowUI
{
private:
	std::unique_ptr<Workflow> mWorkflow;
	WorkflowDatabase* mWorkflowDb;

	ImNodes::EditorContext* mContext;

	ImNodes::NodeId mContextMenuNodeId = 0;
	ImNodes::PinId mContextMenuPinId = 0;
	ImNodes::LinkId mContextMenuLinkId = 0;

public:
	WorkflowUI(std::unique_ptr<Workflow> workflow)
		: mWorkflow{ std::move(workflow) }
	{
		mWorkflowDb = &WorkflowDatabase::GetInstance();
		mContext = ImNodes::CreateEditor();
	}

	~WorkflowUI()
	{
		ImNodes::DestroyEditor(mContext);
	}

	void Draw()
	{
		auto ls = LocaleStrings::Instance.get();

		ImNodes::SetCurrentEditor(mContext);
		ImNodes::Begin("");

		// Defer creation of tooltip because within the node editor, cursor positioning is going to be off
		const char* tooltipMessage = nullptr;

		for (auto& node : mWorkflow->GetNodes()) {
			if (!node) continue;

			ImNodes::BeginNode(node->GetId());
			node->Draw();
			ImNodes::EndNode();
		}

		for (auto& conn : mWorkflow->GetConnections()) {
			if (!conn.IsValid()) continue;

			auto srcId = mWorkflow->GetNodes()[conn.SourceNode]->GetOutputPinUniqueId(conn.SourcePin);
			auto dstId = mWorkflow->GetNodes()[conn.DestinationNode]->GetInputPinUniqueId(conn.DestinationPin);
			ImNodes::Link(conn.GetLinkId(), srcId, dstId);
		}

		if (ImNodes::BeginCreate()) {
			ImNodes::PinId src = 0, dst = 0;
			if (ImNodes::QueryNewLink(&src, &dst)) {
				if (!src || !dst) {
					goto createError;
				}

				auto [srcNode, srcPinId, srcIsOutput] = mWorkflow->DisassembleGlobalPinId(src);
				auto [dstNode, dstPinId, dstIsOutput] = mWorkflow->DisassembleGlobalPinId(dst);

				if (srcNode == dstNode) {
					ImNodes::RejectNewItem();
					goto createError;
				}

				if (srcIsOutput == dstIsOutput) {
					ImNodes::RejectNewItem();
					goto createError;
				}

				auto srcPin = srcNode->GetOutputPin(srcPinId);
				auto dstPin = dstNode->GetOutputPin(dstPinId);

				if (srcPin.MatchingType != dstPin.MatchingType) {
					ImNodes::RejectNewItem();
					goto createError;
				}

				if (ImNodes::AcceptNewItem()) {
					mWorkflow->Connect(*srcNode, srcPinId, *dstNode, dstPinId);
				}
			}

			ImNodes::PinId newNodePin = 0;
			if (ImNodes::QueryNewNode(&newNodePin)) {
				auto [node, pinId, isOutput] = mWorkflow->DisassembleGlobalPinId(newNodePin);

				if ((isOutput && node->GetOutputPin(pinId).IsConnected()) ||
					(!isOutput && node->GetInputPin(pinId).IsConnected()))
				{
					ImNodes::RejectNewItem();
					goto createError;
				}

				if (ImNodes::AcceptNewItem()) {
					ImNodes::Suspend();
					ImGui::BeginPopup("Create Node");
					ImNodes::Resume();
				}
			}
		}
	createError:
		ImNodes::EndCreate();

		if (ImNodes::BeginDelete()) {
			ImNodes::LinkId deletedLinkId;
			if (ImNodes::QueryDeletedLink(&deletedLinkId)) {
				auto& conn = *mWorkflow->GetConnectionByLinkId(deletedLinkId);
				mWorkflow->RemoveConnection(conn.Id);
			}

			ImNodes::NodeId deletedNodeId;
			if (ImNodes::QueryDeletedNode(&deletedNodeId)) {
				auto node = mWorkflow->GetNodeByNodeId(deletedNodeId);
				if (!node) {
					ImNodes::RejectDeletedItem();
					goto deleteError;
				}

				if (node->IsLocked()) {
					ImNodes::RejectDeletedItem();
					goto deleteError;
				}
			}
		}
	deleteError:
		ImNodes::EndDelete();

		// Popups
		ImNodes::Suspend();
		if (ImNodes::ShowNodeContextMenu(&mContextMenuNodeId)) {
			ImGui::OpenPopup("Node Context Menu");
		} else if (ImNodes::ShowPinContextMenu(&mContextMenuPinId)) {
			ImGui::OpenPopup("Pin Context Menu");
		} else if (ImNodes::ShowLinkContextMenu(&mContextMenuLinkId)) {
			ImGui::OpenPopup("Link Context Menu");
		}

		if (ImGui::BeginPopup("Node Context Menu")) {
			auto& node = *mWorkflow->GetNodeByNodeId(mContextMenuNodeId);
			node.DrawDebugInfo();

			if (ImGui::MenuItem(ls->Delete.Get())) {
				ImNodes::DeleteNode(mContextMenuNodeId);
			}

			ImGui::EndPopup();
		}

		if (ImGui::BeginPopup("Pin Context Menu")) {
			auto [node, pinId, isOutput] = mWorkflow->DisassembleGlobalPinId(mContextMenuPinId);
			if (isOutput) {
				node->DrawOutputPinDebugInfo(pinId);
			} else {
				node->DrawInputPinDebugInfo(pinId);
			}

			if (ImGui::MenuItem(ls->Disconnect.Get())) {
				if (isOutput) {
					auto& pin = node->GetOutputPin(pinId);
					if (pin.IsConnected()) {
						auto linkId = mWorkflow->GetConnectionById(pin.Connection)->GetLinkId();
						ImNodes::DeleteLink(linkId);
					}
				} else {
					auto& pin = node->GetInputPin(pinId);
					if (pin.IsConstantConnection()) {
						// TODO
					} else if (pin.IsConnected()) {
						auto linkId = mWorkflow->GetConnectionById(pin.Connection)->GetLinkId();
						ImNodes::DeleteLink(linkId);
					}
				}
			}

			ImGui::EndPopup();
		}

		if (ImGui::BeginPopup("Link Context Menu")) {
			auto& conn = *mWorkflow->GetConnectionByLinkId(mContextMenuLinkId);
			conn.DrawDebugInfo();

			if (ImGui::MenuItem(ls->Delete.Get())) {
				ImNodes::DeleteLink(mContextMenuLinkId);
			}

			ImGui::EndPopup();
		}

		if (ImGui::BeginPopup("Create Node")) {
			auto DisplayCandidatesCategory = [&](const char* name, std::span<const WorkflowDatabase::Candidate> candidates) {
				if (ImGui::BeginMenu(name)) {
					for (auto& candidate : candidates) {
						if (ImGui::MenuItem(candidate.Name.c_str())) {
							// Create node
							auto uptr = candidate.Constructor();
							mWorkflow->AddNode(std::move(uptr));
						}
					}
					ImGui::EndMenu();
				}
			};

			DisplayCandidatesCategory("Numeric nodes", mWorkflowDb->GetNumericNodes());
			DisplayCandidatesCategory("Text nodes", mWorkflowDb->GetTextNodes());
			DisplayCandidatesCategory("Document nodes", mWorkflowDb->GetDocumentNodes());
			DisplayCandidatesCategory("User input nodes", mWorkflowDb->GetUserInputNodes());
			DisplayCandidatesCategory("System input nodes", mWorkflowDb->GetSystemInputNodes());
			DisplayCandidatesCategory("Output nodes", mWorkflowDb->GetOutputNodes());
			ImGui::EndPopup();
		}

		if (tooltipMessage) {
			ImGui::BeginTooltip();
			ImGui::TextUnformatted(tooltipMessage);
			ImGui::EndTooltip();
		}
		ImNodes::Resume();

		ImNodes::End();
	}
};

struct DrawTemplateList_State
{
	const WorkflowInfo* SelectedWorkflow = nullptr;
};

void DrawTemplateList(DrawTemplateList_State& state)
{
	auto& gs = GlobalStates::GetInstance();
	auto& workflows = gs.GetCurrentProject()->GetWorkflows();

	// TODO sort the list
	for (auto& info : workflows) {
		if (ImGui::Selectable(info.Name.c_str(), state.SelectedWorkflow == &info)) {
			state.SelectedWorkflow = &info;
		}
		if (ImGui::IsItemHovered()) {
			ImGui::BeginTooltip();
			ImGui::Text("Path: %s", info.PathStringCache.c_str());
			ImGui::EndTooltip();
		}
	}
}
} // namespace

void UI::WorkflowsTab()
{
	auto ls = LocaleStrings::Instance.get();
	auto& gs = GlobalStates::GetInstance();

	bool openedDummy = true;
	static std::unique_ptr<WorkflowUI> openWorkflow;
	static DrawTemplateList_State state;

	// Toolbar item: close
	if (ImGui::Button(ls->Close.Get(), openWorkflow == nullptr)) {
		openWorkflow = nullptr;
	}

	// Toolbar item: open...
	ImGui::SameLine();
	if (ImGui::Button(ls->OpenWorkflow.Get())) {
		ImGui::OpenPopup(ls->OpenWorkflowDialogTitle.Get());
	}
	if (ImGui::BeginPopupModal(ls->OpenWorkflowDialogTitle.Get(), &openedDummy, ImGuiWindowFlags_AlwaysAutoResize)) {
		DrawTemplateList(state);

		if (state.SelectedWorkflow) {
			auto workflow = state.SelectedWorkflow->LoadFromDisk();
			openWorkflow = std::make_unique<WorkflowUI>(std::move(workflow));
		}

		ImGui::EndPopup();
	}

	// Toolbar item: manage...
	ImGui::SameLine();
	if (ImGui::Button(ls->ManageWorkflows.Get())) {
		ImGui::OpenPopup(ls->ManageWorkflowsDialogTitle.Get());
	}
	if (ImGui::BeginPopupModal(ls->ManageWorkflowsDialogTitle.Get(), &openedDummy, ImGuiWindowFlags_AlwaysAutoResize)) {
		DrawTemplateList(state);

		enum class NameSelectionError
		{
			None,
			Duplicated,
			Empty,
		};
		static std::string newName;
		static NameSelectionError newNameError;
		if (ImGui::Button(ls->Rename.Get(), state.SelectedWorkflow == nullptr)) {
			ImGui::OpenPopup("Rename workflow");
			newName.clear();
		}
		if (ImGui::BeginPopupModal("Rename workflow")) {
			if (ImGui::InputText("New name", &newName)) {
				if (newName.empty()) {
					newNameError = NameSelectionError::Empty;
				}

				auto& workflows = gs.GetCurrentProject()->GetWorkflows();
				if (workflows.find(newName) != workflows.end()) {
					newNameError = NameSelectionError::Duplicated;
				}
			}

			if (ImGui::Button(ls->DialogConfirm.Get(), newName.empty())) {
				auto project = gs.GetCurrentProject();
				project->RenameWorkflow(state.SelectedWorkflow->Name, newName);
				state.SelectedWorkflow = &project->GetWorkflows().at(newName);
			}
			ImGui::SameLine();
			if (ImGui::Button(ls->DialogCancel.Get())) {
				ImGui::CloseCurrentPopup();
			}

			switch (newNameError) {
				case NameSelectionError::None: break;
				case NameSelectionError::Duplicated:
					ImGui::ErrorMessage("Duplicate workflow name");
					break;
				case NameSelectionError::Empty:
					ImGui::ErrorMessage("Workflow name cannot be empty");
					break;
			}

			ImGui::EndPopup();
		}

		if (ImGui::Button(ls->Delete.Get(), state.SelectedWorkflow == nullptr)) {
			ImGui::OpenPopup("Delete confirmation");
		}
		if (ImGui::BeginPopupModal("Delete confirmation")) {
			if (ImGui::Button(ls->DialogConfirm.Get())) {
				gs.GetCurrentProject()->RemoveWorkflow(state.SelectedWorkflow->Name);
			}
			ImGui::SameLine();
			if (ImGui::Button(ls->DialogCancel.Get())) {
				ImGui::CloseCurrentPopup();
			}
			ImGui::EndPopup();
		}

		ImGui::EndPopup();
	}

	if (openWorkflow) {
		openWorkflow->Draw();
	}
}