#include <cstdio>
#include <cstdlib>
#include <cassert>

#include <iostream>
#include <sstream>
#include <fstream>

#include <map>
#include <utility>
#include <string>
#include <set>

#include <vector>
#include <algorithm>

using namespace std;

#define OUT cout

typedef vector<string> exp_path_t;
typedef vector<exp_path_t> exp_path_list_t;

void read_exp_path(const string & line, exp_path_t & path)
{
	istringstream iss(line);
	string token;

	iss >> token;
	assert(token == "[");
	
	path.clear();
	while (true) {
		assert(iss >> token);
		if (token == "]")
			break;
		path.push_back(token);
	}
}

void read_exp_path_list(char * filename, exp_path_list_t & path_list)
{
	string line;
	ifstream ifs(filename);
	assert(ifs.is_open());

	path_list.clear();
	while (getline(ifs, line)) {
		if (line == "goal")
			break;
	} 
	assert(line == "goal");

	while (getline(ifs, line)) {
		if (line == "")
			break;
		path_list.push_back(exp_path_t());
		read_exp_path(line, path_list.back());
	}
}

void usage(char * exe_file)
{
	/* (1) print the list of exploit paths from a file
	 * (2) compare two lists of exploit paths
	 */
	cout << "usage (1): " << exe_file << " list listfile" << endl;
	cout << "usage (2): " << exe_file << " compare listfile1 listfile2" << endl;

	exit(0);
}

void check_cmd(int argc, char ** argv)
{
	if (argc < 2)
		usage(argv[0]);

	if (strcmp(argv[1], "list") == 0)
		assert(argc == 3);
	else if (strcmp(argv[1], "compare") == 0)
		assert(argc == 4);
	else {
		cerr << "unknown option: '" << argv[1] << "'" << endl;
		usage(argv[0]);
	}
}

string tex_format(const string & str)
{
	string result;
	for (size_t i = 0; i < str.size(); ++i) {
		if (str[i] == '_')
			result += '\\';
		result += str[i];
	}
	return result;
}

void do_list(char * filename)
{
	map<int, int> surface_sig;
	int max_sig_index = 0;

	exp_path_list_t path_list;
	read_exp_path_list(filename, path_list);

	//output the list in a latex table with numerical labels
	OUT << "\\begin{tabular}{|c|p{7cm}|}" << endl;
	OUT << "\\hline" << endl;
	OUT << "index & exploit path \\\\" << endl;
	OUT << "\\hline" << endl;

	for (size_t i = 0; i < path_list.size(); ++i) {
		OUT << i+1 << " & ";
		for (size_t j = 0; j < path_list[i].size(); ++j)
			OUT << tex_format(path_list[i][j]) << " ";
		OUT << "\\\\" << endl;
		OUT << "\\hline" << endl;

		++surface_sig[path_list[i].size()];
		max_sig_index = max(max_sig_index, int(path_list[i].size()));
	}

	OUT << "\\end{tabular}" << endl;

	OUT << endl << "$\\left\\langle ";
	for (int i = 1; i <= max_sig_index; ++i)
		OUT << surface_sig[i] << ", ";
	OUT << "0 \\dots \\right\\rangle$" << endl;
}

//if a path is 'stronger' than another to the attacker
//A be stronger than B means A is a subset of B
bool stronger_path(const exp_path_t & exp_path1, const exp_path_t & exp_path2)
{
	for (size_t i = 0; i < exp_path1.size(); ++i) {
		bool found = false;
		for (size_t j = 0; j < exp_path2.size(); ++j)
			if (exp_path1[i] == exp_path2[j]) {
				found = true;
				break;
			}
		if (!found)
			return false;
	}
	return true;
}

void categorize(const exp_path_list_t & path_list1,
				const exp_path_list_t & path_list2,
				set<int> & weak,
				set<int> & common,
				set<int> & strong,
				set<int> & unique)
{
	weak.clear();
	common.clear();
	strong.clear();
	unique.clear();

	for (size_t i = 0; i < path_list1.size(); ++i) {
		bool found = false;

		for (size_t j = 0; j < path_list2.size(); ++j) {
			bool st12 = stronger_path(path_list1[i], path_list2[j]);
			bool st21 = stronger_path(path_list2[j], path_list1[i]);
			if (st12 && st21)
				common.insert(i);
			if (st12 && !st21)
				strong.insert(i);
			if (!st12 && st21)
				weak.insert(i);
			if (st12 || st21) {
				found = true;
				break;
			}
		}

		if (!found)
			unique.insert(i);
	}
}

void output_compare_row(const string & head,
						const set<int> & ids1, const set<int> & ids2)
{
	OUT << head << " &";
	for (typeof(ids1.begin()) it = ids1.begin(); it != ids1.end(); ++it)
		OUT << " " << *it + 1;
	OUT << " &";
	for (typeof(ids2.begin()) it = ids2.begin(); it != ids2.end(); ++it)
		OUT << " " << *it + 1;
	OUT << " \\\\" << endl;
	OUT << "\\hline" << endl;
}

void output_path_list(const string & head,
						const exp_path_list_t & path_list,
						const set<int> & ids)
{
	OUT << endl << head << ":" << endl;
	OUT << "==============================" << endl;
	for (typeof(ids.begin()) it = ids.begin(); it != ids.end(); ++it) {
		for (int i = 0; i < int(path_list[*it].size()); ++i)
			OUT << "  " << path_list[*it][i];
		OUT << endl;
	}
}

void do_compare(char * filename1, char * filename2)
{
	exp_path_list_t path_list[2];
	set<int> weak[2];
	set<int> common[2];
	set<int> strong[2];
	set<int> unique[2];

	read_exp_path_list(filename1, path_list[0]);
	read_exp_path_list(filename2, path_list[1]);

	for (int i = 0; i < 2; ++i)
		categorize(path_list[i], path_list[1-i],
					weak[i], common[i], strong[i], unique[i]);

	//output the result as a table
	OUT << "\\begin{tabular}{|l|p{3cm}|p{3cm}|}" << endl;
	OUT << "\\hline" << endl;
	OUT << "& " << filename1 << " & " << filename2 << "\\\\" << endl;
	OUT << "\\hline" << endl;

	output_compare_row("uniq", unique[0], unique[1]);
	output_compare_row("strong", strong[0], strong[1]);
	output_compare_row("common", common[0], common[1]);
	output_compare_row("weak", weak[0], weak[1]);

	OUT << "\\end{tabular}" << endl;

	//output a table with details
	output_path_list("unique", path_list[0], unique[0]);
	output_path_list("strong", path_list[0], strong[0]);
	output_path_list("common", path_list[0], common[0]);
	output_path_list("weak", path_list[0], weak[0]);

	//output a simple table
	OUT << "\\begin{tabular}{|l|p{3cm}|p{3cm}|}" << endl;
	OUT << "\\hline" << endl;
	OUT << "& " << filename1 << " & " << filename2 << "\\\\" << endl;
	OUT << "uniq   & " << unique[0].size() << " & " <<  unique[1].size() << "\\\\" << endl;
	OUT << "strong & " << strong[0].size() << " & " <<  strong[1].size() << "\\\\" << endl;
	OUT << "common   & " << common[0].size() << " & " <<  common[1].size() << "\\\\" << endl;
	OUT << "weak   & " << weak[0].size() << " & " <<  weak[1].size() << "\\\\" << endl;
	OUT << "\\end{tabular}" << endl;
}

int main(int argc, char ** argv)
{
	check_cmd(argc, argv);

	if (strcmp(argv[1], "list") == 0)
		do_list(argv[2]);
	else if (strcmp(argv[1], "compare") == 0)
		do_compare(argv[2], argv[3]);

	return 0;
}

