#include <cctype>
#include <cstdlib>
#include <functional>
#include <iostream>
#include <fstream>
#include <optional>
#include <ostream>
#include <random>
#include <string>
#include <sstream>
#include <vector>

#include "pcg_random.hpp"
#include "eff_large_wordlist.txt.hpp"
#include "eff_short_wordlist_1.txt.hpp"
#include "diceware-wordlist-8k-composites-nl.txt.hpp"

enum WordlistType { WORDLIST_EFF_LARGE, WORDLIST_EFF_SHORT, WORDLIST_DUTCH, WORDLIST_CUSTOM };

struct PassgenOptions {
	bool help = false;
	WordlistType wordlist_type = WORDLIST_EFF_SHORT;
	int wordcount = 5;
	bool camelcase = false;
	std::string seperator = "-";
	std::optional<std::string> wordlist_path;
};

typedef std::function<void(PassgenOptions&)> NoArgHandle;
typedef std::function<void(PassgenOptions&, const std::string&)> OneArgHandle;

const std::unordered_map<std::string, NoArgHandle> NoArgs {
	{"-h", [](PassgenOptions& s) { s.help = true; }},
	{"--help", [](PassgenOptions& s) { s.help = true; }},

	{"-nl", [](PassgenOptions& s) { s.wordlist_type = WORDLIST_DUTCH; }},
	{"--dutch", [](PassgenOptions& s) { s.wordlist_type = WORDLIST_DUTCH; }},

	{"-s", [](PassgenOptions& s) { s.wordlist_type = WORDLIST_EFF_SHORT; }},
	{"--short", [](PassgenOptions& s) { s.wordlist_type = WORDLIST_EFF_SHORT; }},

	{"-l", [](PassgenOptions& s) { s.wordlist_type = WORDLIST_EFF_LARGE; }},
	{"--large", [](PassgenOptions& s) { s.wordlist_type = WORDLIST_EFF_LARGE; }},

	{"--camelcase", [](PassgenOptions& s) { s.camelcase = true; }},
	{"-cc", [](PassgenOptions& s) { s.camelcase = true; }},
};

const std::unordered_map<std::string, OneArgHandle> OneArgs {
	{"-wl", [](PassgenOptions& s, const std::string& arg) {
		s.wordlist_path = arg;
		s.wordlist_type = WORDLIST_CUSTOM;
	}},
	{"--wordlist", [](PassgenOptions& s, const std::string& arg) {
		s.wordlist_path = arg;
		s.wordlist_type = WORDLIST_CUSTOM;
	}},
	{"-n", [](PassgenOptions& s, const std::string& arg) {
		s.wordcount = std::stoi(arg);
	}},
	{"-wc", [](PassgenOptions& s, const std::string& arg) {
		s.wordcount = std::stoi(arg);
	}},
	{"--wordcount", [](PassgenOptions& s, const std::string& arg) {
		s.wordcount = std::stoi(arg);
	}},
	{"--seperator", [](PassgenOptions& s, const std::string& arg) {
		s.seperator = arg;
	}},
};

PassgenOptions parse_options(int argc, const char* argv[]) {
	PassgenOptions settings;

	for(int i = 1; i < argc; i++) {
		std::string opt = argv[i];

		if (auto j = NoArgs.find(opt); j != NoArgs.end()) {
			j->second(settings);
		} else if(auto k = OneArgs.find(opt); k != OneArgs.end())  {
			if(++i < argc) {
				k->second(settings, {argv[i]});
			} else {
				throw std::runtime_error {"missing param after " + opt};
				std::exit(1);
			}
		} else {
			std::cerr << "unrecognized command-line option " << opt << std::endl;
			std::exit(1);
		}
	}
	return settings;
}

std::vector<std::string> getWordlist(PassgenOptions opts) 
{
	std::vector<std::string> map;

	if (opts.wordlist_type != WORDLIST_CUSTOM)
	{
		std::stringstream ss;
		if (opts.wordlist_type == WORDLIST_EFF_SHORT) {
			ss.write(reinterpret_cast<const char*>(eff_short_wordlist_1_txt), eff_short_wordlist_1_txt_len);
		} else if (opts.wordlist_type == WORDLIST_EFF_LARGE){
			ss.write(reinterpret_cast<const char*>(eff_large_wordlist_txt), eff_large_wordlist_txt_len);
		} else if (opts.wordlist_type == WORDLIST_DUTCH){
			ss.write(reinterpret_cast<const char*>(diceware_wordlist_8k_composites_nl_txt), diceware_wordlist_8k_composites_nl_txt_len);
		}
		std::string line;	
		while (std::getline(ss, line)) {
			map.push_back(line);
		}
	}
	else if (opts.wordlist_path.has_value())
	{
		std::string path = opts.wordlist_path.value();
		std::ifstream file(path);
		if (file.is_open()) {
			std::string line;	
			while (std::getline(file, line)) {
				map.push_back(line);
			}
		} else {
			std::cout << "ERROR: failed to open '" << path << "'" << std::endl;
		}
	} else {
			std::cout << "ERROR: no wordlist specified" << std::endl;
	}

	return map;
}

int main(int argc, const char *argv[]) {
	PassgenOptions opts = parse_options(argc, argv);

	if (opts.help) {
		std::cout << R"(
Usage: passgen [options]

Options:
  -h, --help              Display this help message and exit
  -s, --short             Use the short EFF wordlist (dfault)
  -l, --large             Use the large EFF wordlist
  --camelcase, -cc        Generate passphrase in CamelCase format			
  -wl <path>, --wordlist <path>
                          Use a custom wordlist file located at <path>
  -n <count>, -wc <count>, --wordcount <count>
                          Amount of words to use for passphrase, default = 5
  --separator <seperator> Specify a character to use as a separator between words, default = `-`

Description:
  Generate passphrases based on wordlists.

Examples:
  passgen -h
  passgen --short 4
  passgen -wl /path/to/custom_wordlist.txt --wordcount 4
)";
		return 0;
	}

	std::vector<std::string> wordlist = getWordlist(opts);

	if (wordlist.size() == 0) {
		std::cerr << "Wordlist is empty";
		return 1;	
	}

	pcg_extras::seed_seq_from<std::random_device> seed_source;
	pcg32 rng(seed_source);

	for (int i = 0; i < opts.wordcount; i++) {
		int random_int = std::uniform_int_distribution<int>(0, wordlist.size())(rng);

		if (opts.camelcase) {
			wordlist[random_int][0] = toupper(wordlist[random_int][0]);
		}
		std::cout << wordlist[random_int];

		if (i+1 == opts.wordcount) {
			std::cout << "\n";
		} else {
			std::cout << opts.seperator;
		}
	}
}