#include "aes.hpp"

#include <iostream>
#include <fstream>
#include <string>

#include "profiler.hpp"

using namespace std;

uint8_t key[16];

string prompt(char separator = '>')
{
	PROFILER_RECORD;
	string input;
	cout << separator << ' ';
	cin >> input;
	if (!cin.good())
	{
		cout << "OOPSIE WOOPSIE!! Uwu We made a fucky wucky!! A wittle fucko boingo! The code monkeys at our headquarters are working VEWY HAWD to fix this!" << endl;
		exit(1);
	}
	return input;
}

void print_help()
{
	PROFILER_RECORD;
	// TODO Verify whether there are more/less commands
	cout << "The following commands are available: help, login, register, exit" << endl;
}

void print_flag()
{
	PROFILER_RECORD;
	ifstream file("flag.txt");
	string flag;
	getline(file, flag);
	cout << flag << endl;
	exit(0);
}

void login()
{
	PROFILER_RECORD;
	cout << "Username";
	string username = prompt(':');
	cout << "Password";
	string password = prompt(':');
	if (username != "admin")
	{
		cout << "Unknown user '" << username << "'" << endl;
		return;
	}
	if (password.length() != 16)
	{
		cout << "Password format error" << endl;
		return;
	}
	uint8_t encrypted_password[16];
	for (int i = 0;i < 16;i++)
	{
		encrypted_password[i] = password[i];
	}
	AES::encrypt_ecb(encrypted_password, key);
	uint8_t rootpw[16] = {0xe3, 0x07, 0x2e, 0x9f, 0x5b, 0xe8, 0xed, 0xd6, 0x02, 0xab, 0x89, 0xb8, 0xeb, 0x49, 0xcc, 0x56};
	for (int i = 0;i < 16;i++)
	{
		if (encrypted_password[i] != rootpw[i])
		{
			cout << "Invalid password" << endl;
			return;
		}
	}
	print_flag();
}

#include <iomanip>
void register_user()
{
	PROFILER_RECORD;
	cout << "Registration is currently unavailable" << endl;
}

bool execute_command(const string &input)
{
	PROFILER_RECORD;
	if (input == "exit")
	{
		cout << "Thank you for visiting the Secutech customer portal. Please visit us again soon!" << endl;
		return false;
	}
	if (input == "login")
	{
		login();
	}
	else if (input == "register")
	{
		register_user();
	}
	else if (input == "help")
	{
		print_help();
	}
	else if (input == "profiler_print")
	{
		profiler.print();
	}
	else if (input == "profiler_reset")
	{
		profiler.reset();
	}
	else
	{
		cout << "Invalid command '" << input << "'" << endl;
	}
	return true;
}

void load_key()
{
	PROFILER_RECORD;
	ifstream file("key.txt");
	string line;
	getline(file, line);
	for (int i = 0;i < 16;i++)
	{
		key[i] = stoi(line.substr(3 * i, 2), 0, 16);
	}
}

int main(int argc, char **argv) {
	PROFILER_RECORD;
	load_key();
	cout << "Welcome to the Secutech customer portal." << endl;
	cout << "How can we help you today? (type 'help' for help)" << endl;
	string input;
	do
	{
		input = prompt();
	} while (execute_command(input));

	return 0;
}