#include "totpgen.h"

#ifdef _WIN32
static const std::string path_separator = "\\";
static const std::string config_dir_name = "TOTPGen";
#else
static const std::string path_separator = "/";
static const std::string config_dir_name = ".totpgen";
#endif

IMPLEMENT_APP(TOTPGenGUI)

bool
TOTPGenGUI::OnInit()
{
	TOTPGen *MainAppFrame = new TOTPGen();
	MainAppFrame->Show(true);
	return true;
}

TOTPGen::TOTPGen() : wxFrame(nullptr, wxID_ANY, wxT("TOTP Generator"),
                             wxDefaultPosition, wxSize(320,320),
                             wxDEFAULT_FRAME_STYLE & ~(wxRESIZE_BORDER | wxMAXIMIZE_BOX))
{
	MainWindow = new wxPanel(this, wxID_ANY);
	MainSizer = new wxBoxSizer(wxVERTICAL);

	KeyDB = {};

	wxString UserConfigDir = wxStandardPaths::Get().GetUserConfigDir();

	wxString AppConfigDir = UserConfigDir + path_separator + config_dir_name;
	if (!wxDirExists(AppConfigDir)) {
		wxMkdir(AppConfigDir);
	}

	SecretDB_File = AppConfigDir + path_separator + "secret.db";

	SecretChooser = new wxComboBox(MainWindow, wxID_ANY, wxEmptyString,
	                               wxDefaultPosition, wxDefaultSize, 0, NULL,
	                               wxCB_READONLY, wxDefaultValidator);
	SecretChooser->Enable(false);
	SecretChooser->Bind(wxEVT_COMBOBOX, &TOTPGen::SelectSecret, this);

	SelectSizer = new wxBoxSizer(wxHORIZONTAL);
	SelectSizer->Add(SecretChooser, 2,  wxEXPAND | wxALL);

	TOTPField = new wxStaticText(MainWindow, wxID_ANY,
	                             wxEmptyString, wxDefaultPosition,
	                             wxDefaultSize, wxTE_CENTER);
	wxFont LabelFont(48, wxFONTFAMILY_TELETYPE,
	                     wxFONTSTYLE_NORMAL,
	                     wxFONTWEIGHT_NORMAL);
	TOTPField->SetFont(LabelFont);
	TOTPField->SetLabel("      ");
	ResizeTextToFit(TOTPField);

	TextSizer = new wxBoxSizer(wxHORIZONTAL);
	TextSizer->Add(TOTPField, 2, wxCENTER);

	CountdownGauge = new wxGauge(MainWindow, wxID_ANY, 30, wxPoint(50, 50));
	CountdownTimer = new wxTimer(MainWindow, wxID_ANY);

	MainWindow->Connect(CountdownTimer->GetId(), wxEVT_TIMER,
	    wxTimerEventHandler(TOTPGen::UpdateStatus), NULL, this);

	AddButton = new wxButton(MainWindow, wxID_OK, wxT("Add"));
	AddButton->Bind(wxEVT_BUTTON, &TOTPGen::AddSecret, this);
	AddButton->Enable(false);

	DeleteButton = new wxButton(MainWindow, wxID_OK, wxT("Delete"));
	DeleteButton->Bind(wxEVT_BUTTON, &TOTPGen::DeleteSecret, this);
	DeleteButton->Enable(false);

	ButtonSizer_Upper = new wxBoxSizer(wxHORIZONTAL);
	ButtonSizer_Upper->Add(AddButton, 2, wxEXPAND | wxALL);
	ButtonSizer_Upper->Add(DeleteButton, 2, wxEXPAND | wxALL);

	CopyButton = new wxButton(MainWindow, wxID_COPY, wxT("Copy"));
	CopyButton->Bind(wxEVT_BUTTON, &TOTPGen::CopyText, this);
	CopyButton->Enable(false);

	QuitButton = new wxButton(MainWindow, wxID_EXIT, wxT("Exit"));
	QuitButton->Bind(wxEVT_BUTTON, &TOTPGen::OnExit, this);

	ButtonSizer_Middle = new wxBoxSizer(wxHORIZONTAL);
	ButtonSizer_Middle->Add(CopyButton, 2, wxEXPAND | wxALL);
	ButtonSizer_Middle->Add(QuitButton, 2, wxEXPAND | wxALL);

	SecurityButton = new wxButton(MainWindow, wxID_OK, wxT("Unlock Database"));
	SecurityButton->Bind(wxEVT_BUTTON, &TOTPGen::Unlock, this);
	if (!wxFileExists(SecretDB_File)) {
		SecurityButton->SetLabel(wxT("Create Database"));
		SecurityButton->Bind(wxEVT_BUTTON, &TOTPGen::InitDB, this);
	}

	ButtonSizer_Lower = new wxBoxSizer(wxHORIZONTAL);
	ButtonSizer_Lower->Add(SecurityButton, 2, wxEXPAND | wxALL);

	MainSizer->Add(SelectSizer, 2, wxEXPAND | wxALL);
	MainSizer->Add(TextSizer, 2, wxCENTER | wxALIGN_CENTER);
	MainSizer->Add(CountdownGauge, 2, wxEXPAND | wxALL);
	MainSizer->Add(ButtonSizer_Upper, 2, wxEXPAND | wxALL);
	MainSizer->Add(ButtonSizer_Middle, 2, wxEXPAND | wxALL);
	MainSizer->Add(ButtonSizer_Lower, 2, wxEXPAND | wxALL);

	MainWindow->SetSizer(MainSizer);
}

void
TOTPGen::OnExit(wxCommandEvent& WXUNUSED(event))
{
	Close(true);
}

bool
TOTPGen::IsBase32(std::string& str)
{
	std::unordered_set<char> Base32Chars(base32.begin(), base32.end());

	for (char c: str) {
		if (Base32Chars.find(c) == Base32Chars.end()) {
			return false;
		}
	}

	return true;
}

void
TOTPGen::ReloadSelections(void)
{
	SecretChooser->Clear();

	int pos = 0;
	for (auto Entry : KeyDB.items()) {
		SecretChooser->Insert(Entry.key(), pos);
		pos++;
	}

	if (pos) {
		DeleteButton->Enable(true);
		CopyButton->Enable(true);

		SecretChooser->Enable(true);
		SecretChooser->SetSelection(0);

		wxCommandEvent DummyComboEvent(wxEVT_COMBOBOX);
		SelectSecret(DummyComboEvent);

		CountdownTimer->Start(1000);
	} else {
		DeleteButton->Enable(false);
		CopyButton->Enable(false);
		SecretChooser->Enable(false);

		CountdownTimer->Stop();
		CountdownGauge->SetValue(0);

		TOTPField->SetLabel("      ");
	}
}

void
TOTPGen::SelectSecret(wxCommandEvent& WXUNUSED(event))
{
	Secret = KeyDB[SecretChooser->GetStringSelection().ToStdString()];
	GenTOTP(Secret);
}

void
TOTPGen::DeleteSecret(wxCommandEvent& WXUNUSED(event))
{
	wxString SecretName_Selected = SecretChooser->GetStringSelection();

	wxMessageDialog *ConfirmDialog = new wxMessageDialog(this,
	    wxT("Are you want to delete the token '" + SecretName_Selected + "'?"),
	    wxT("Confirm Deletion"), wxYES_NO | wxNO_DEFAULT | wxICON_QUESTION);
	int ConfirmResponse = ConfirmDialog->ShowModal();

	if (ConfirmResponse != wxID_YES) {
		return;
	}

	KeyDB.erase(SecretName_Selected.ToStdString());

	EncryptDB(KeyDB.dump().c_str());

	ReloadSelections();
}

void
TOTPGen::AddSecret(wxCommandEvent& WXUNUSED(event))
{
	struct {
		std::string Name;
		std::string Base32Str;
	} NewSecret;

	wxTextEntryDialog SecretPrompt(this, wxT("Please enter the base 32 token."),
	                               wxT("Secret Entry"), wxEmptyString);

	if (SecretPrompt.ShowModal() == wxID_OK) {
		NewSecret.Base32Str = SecretPrompt.GetValue().ToStdString();
		if (!IsBase32(NewSecret.Base32Str) || NewSecret.Base32Str.empty()) {
			wxMessageBox(wxT("Base 32 token contains invalid characters.\n"
			                 "Only the following are permissible: \n" + base32),
			             wxT("Error"), wxOK | wxICON_ERROR);
			return;
		}
	} else {
		return;
	}

	wxTextEntryDialog NamePrompt(this, wxT("Please enter a name for the secret."),
	                             wxT("Secret Entry"), wxEmptyString);

	if (NamePrompt.ShowModal() == wxID_OK) {
		NewSecret.Name = NamePrompt.GetValue().ToStdString();
	} else {
		return;
	}

	if (NewSecret.Base32Str.empty() || NewSecret.Name.empty()) {
		wxMessageBox(wxT("Empty input - please try again."),
		             wxT("Error"), wxOK | wxICON_ERROR);
		return;
	}

	if (!GenTOTP(NewSecret.Base32Str)) {
		wxMessageBox(wxT("Could not generate TOTP - invalid token?"),
		             wxT("Error"), wxOK | wxICON_ERROR);
		return;
	}

	KeyDB[NewSecret.Name] = NewSecret.Base32Str;

	EncryptDB(KeyDB.dump().c_str());

	ReloadSelections();
}

void
TOTPGen::InitDB(wxCommandEvent& WXUNUSED(event))
{
	wxCommandEvent DummyBtnEvent(wxEVT_BUTTON);
	SetPassword(DummyBtnEvent);

	AddButton->Enable(true);

	SecurityButton->SetLabel(wxT("Set Password"));
	SecurityButton->Bind(wxEVT_BUTTON, &TOTPGen::SetPassword, this);
}

void
TOTPGen::SetPassword(wxCommandEvent& WXUNUSED(event))
{
	wxPasswordEntryDialog PasswordPrompt(this, wxT("Please enter a new passphrase\n"
	                                               "to encrypt the token database."),
	                                     wxT("Passphrase Entry"), wxEmptyString);

	wxPasswordEntryDialog ConfirmPrompt(this, wxT("Please confirm the passphrase."),
	                                     wxT("Passphrase Entry"), wxEmptyString);

	if ((PasswordPrompt.ShowModal() == wxID_OK)
	     && (ConfirmPrompt.ShowModal() == wxID_OK)) {
		if (PasswordPrompt.GetValue() == ConfirmPrompt.GetValue()) {
			Passphrase = PasswordPrompt.GetValue();
		}
		else {
			wxMessageBox(wxT("Passphrases did not match - try again"),
			             wxT("Error"), wxOK | wxICON_ERROR);
			return;
		}
		if (Passphrase.empty()) {
			wxMessageBox(wxT("Passphrase may not be empty."),
			wxT("Error"), wxOK | wxICON_ERROR);
			return;
		}
	} else {
		return;
	}

	EncryptDB(KeyDB.dump().c_str());
}

void
TOTPGen::Unlock(wxCommandEvent& WXUNUSED(event))
{
	wxPasswordEntryDialog PasswordPrompt(this, wxT("Please enter your passphrase."),
	                                     wxT("Passphrase Entry"), wxEmptyString);

	if (PasswordPrompt.ShowModal() == wxID_OK) {
		Passphrase = PasswordPrompt.GetValue();
	}

	if (!DecryptDB()) {
		wxMessageBox(wxT("Could not unlock database,\n"
		                 "bad passphrase or corrupt file?"),
		             wxT("Error"), wxOK | wxICON_ERROR);
		return;
	}

	AddButton->Enable(true);
	SecurityButton->SetLabel(wxT("Set Password"));
	SecurityButton->Bind(wxEVT_BUTTON, &TOTPGen::SetPassword, this);

	ReloadSelections();
	GenTOTP(Secret);
}

void
TOTPGen::UpdateStatus(wxTimerEvent& WXUNUSED(event))
{
	unsigned char CounterValue = 29 - ((time(NULL) - 1) % 30);

	CountdownGauge->SetValue(CounterValue);

	if (CounterValue == 0) {
		GenTOTP(Secret);
	}
}

void
TOTPGen::CopyText(wxCommandEvent& WXUNUSED(event))
{
	wxString ToCopy = TOTPField->GetLabel();

	if (wxTheClipboard->Open()) {
		wxTheClipboard->SetData(new wxTextDataObject(ToCopy));
		wxTheClipboard->Close();
	}
}

void
TOTPGen::ResizeTextToFit(wxStaticText* text_field)
{
	wxString text = text_field->GetLabel();
	wxSize text_size = text_field->GetClientSize();
#if defined(__WXGTK__)
	double GTK_scale = std::stod(getenv("GDK_DPI_SCALE"));
	text_size.x /= GTK_scale;
	text_size.y /= GTK_scale;
#endif

	wxFont font = text_field->GetFont();
	int font_size = font.GetPointSize();

	wxClientDC device_ctx(text_field);
	device_ctx.SetFont(font);

	int text_width, text_height;
	device_ctx.GetTextExtent(text, &text_width, &text_height);

	while (text_width > (text_size.x - 12)
		     || text_height > (text_size.y - 12)) {
		font_size--;
		font.SetPointSize(font_size);
		device_ctx.SetFont(font);
		device_ctx.GetTextExtent(text, &text_width, &text_height);
	}

	text_field->SetFont(font);
	text_field->Refresh();
}