#include <string>
#include <unistd.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <pwd.h>
#include <unistd.h>
#include <iostream>
#include <csignal>
#include <fstream>
#include <vector>
#include <ncurses.h>
#include <time.h>
#include <stdlib.h>
#include <stdio.h>
#include "sha512.h"

#define WAIT_TICK 10000

using namespace std;

WINDOW *wnd;

string pw="";
char pwchar='*';
bool randomchar=true;
bool exit_requested=false;
bool changepassword=false;
bool softlock=false;

int pwfieldlength=40;

int LINES;
int COLS;

vector<string> quitmsgs;
vector<string> wrongmsgs;

string pwpath="";
string pwhash="3c9909afec25354d551dae21590bb26e38d53f2173b8d3dc3eee4c047e7ab1c1eb8b85103e3be7ba613b31bb5c9c36214dc9f14a42fd7a2fdb84856bca5c44c2";

int init() {
	srand(time(NULL));
	fstream f(pwpath, ios::in);
	if (f.good()) {
		getline(f, pwhash);
		if (pwhash.length()<20) {
			cerr << "Error: passwd file not correctly formatted." << endl;
			exit(1);
		}
	} else {
		cerr << "Error while loading password file." << endl;
		exit(1);
	}
	f.close();

	initscr();
	wnd=newwin(0,0,0,0);
	
	getmaxyx(wnd,LINES,COLS);

	noecho();
	clear();

	keypad(wnd, true);

	//curs_set(0);

	start_color();

	init_pair(1,COLOR_WHITE, COLOR_BLACK);
	init_pair(2,COLOR_WHITE, COLOR_BLUE);
	init_pair(3,COLOR_RED, COLOR_BLACK);

	bkgd(COLOR_PAIR(1));

	string enterpw="Enter password:";
	for (int i=0; i<enterpw.length(); i++) {
		mvaddch(LINES/2-2,COLS/2-enterpw.length()/2+i,enterpw[i]);
	}
	
	for (int i=0; i<pwfieldlength; i++) {
		mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,' ' | COLOR_PAIR(2));
	}
	move(LINES/2,COLS/2-pwfieldlength/2);
	wrefresh(wnd);

	return 0;
}

void handler(int signum)
{
	if (changepassword || softlock) {
		echo();
		endwin();
		exit(1);
	} else {
		for (int i=0; i<COLS; i++) {
			mvaddch(LINES/2+2,i,' ');
		}
		int printindex=rand() % quitmsgs.size();
		string str=quitmsgs[printindex];
		attron(A_BOLD);
		for (int i=0; i<str.length(); i++) {
			mvaddch(LINES/2 + 2, COLS/2 - str.length()/2+i, str[i] | COLOR_PAIR(3));
		}
		attroff(A_BOLD);
		if (pw.length()<pwfieldlength) {
			move(LINES/2,COLS/2-pwfieldlength/2+pw.length());
		} else {
			move(LINES/2,COLS/2+pwfieldlength/2);
		}
		refresh();
	}
}

void checkpw() {
	string crypt=sha512(pw);
	if (crypt!= pwhash) {
		if (!changepassword && !softlock) {
			usleep(1000000);
			for (int i=0; i<COLS; i++) {
				mvaddch(LINES/2+2,i,' ');
			}
			int printindex=rand() % wrongmsgs.size();
			string str=wrongmsgs[printindex];
			attron(A_BOLD);
			for (int i=0; i<str.length(); i++) {
				mvaddch(LINES/2+2,COLS/2-str.length()/2+i,str[i] | COLOR_PAIR(3));
			}
			attroff(A_BOLD);
			if (pw.length()<pwfieldlength) {
				move(LINES/2,COLS/2-pwfieldlength/2+pw.length());
			} else {
				move(LINES/2,COLS/2+pwfieldlength/2);
			}
			refresh();
			pw="";
		} else {
			usleep(1000000);
			echo();
			endwin();
			cerr << "Wrong password!" << endl;
			exit(1);
		}
	} else {
		exit_requested=true;
		endwin();
		if (!changepassword) {
			exit(0);
		}
	}
}

void run() {
	while(!exit_requested) {
		if (cin.eof()) {}

		int in_char=getch();

		for (int i=0; i<COLS; i++) {
			mvaddch(LINES/2+2,i,' ');
		}
		if (in_char==263 || in_char==330 || in_char=='\b' || in_char==127) {  //  BACKSPACE
			pw=pw.substr(0,pw.size()-1);
		} else if (in_char == 10 || in_char=='\n') {             //ENTER
			checkpw();
		} else {                               //LETTERS
			pw = pw + (char)in_char;
		}
		for (int i=0; i<pwfieldlength; i++) {
			if (i<pw.length()) {
				if (randomchar) {
					char ch = (char) (rand() % 93 + 33);
					mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,ch | COLOR_PAIR(2));
				} else {
					mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,pwchar | COLOR_PAIR(2));
				}
			} else {
				mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,' ' | COLOR_PAIR(2));
			}
		}
		if (pw.length()<pwfieldlength) {
			move(LINES/2,COLS/2-pwfieldlength/2+pw.length());
		} else {
			move(LINES/2,COLS/2+pwfieldlength/2);
		}
		usleep(WAIT_TICK);
		refresh();
	}
	checkpw();
}


string readPw(string msg) {
	initscr();
	wnd=newwin(0,0,0,0);
	getmaxyx(wnd,LINES,COLS);

	noecho();
	clear();

	keypad(wnd, true);

	//curs_set(0);

	start_color();

	init_pair(1,COLOR_WHITE, COLOR_BLACK);
	init_pair(2,COLOR_WHITE, COLOR_BLUE);

	pw="";
	bkgd(COLOR_PAIR(1));
	exit_requested=false;

	for (int i=0; i<msg.length(); i++) {
		mvaddch(LINES/2-2, COLS/2-msg.length()/2+i,msg[i]);
	}
	for (int i=0; i<pwfieldlength; i++) {
		if (i<pw.length()) {
			if (randomchar) {
				char ch = (char)(rand() % 93 + 33);
				mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,ch | COLOR_PAIR(2));
			} else {
				mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,pwchar | COLOR_PAIR(2));
			}
		} else {
			mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,' ' | COLOR_PAIR(2));
		}
	}
	for (int i=0; i<COLS; i++) {
		mvaddch(LINES/2+2,i,' ');
	}
	if (pw.length()<pwfieldlength) {
		move(LINES/2,COLS/2-pwfieldlength/2+pw.length());
	} else {
		move(LINES/2,COLS/2+pwfieldlength/2);
	}
	refresh();
	while(!exit_requested) {
		if (cin.eof()) {}

		int in_char=getch();

		if (in_char==263 || in_char==330 || in_char=='\b' || in_char==127) {  //  BACKSPACE
			pw=pw.substr(0,pw.size()-1);
		} else if (in_char == 10 || in_char=='\n') {             //ENTER
			echo();
			endwin();
			return pw;
		} else {                               //LETTERS
			pw= pw + (char)in_char;
		}
		for (int i=0; i<pwfieldlength; i++) {
			if (i<pw.length()) {
				if (randomchar) {
					char ch = (char)(rand() % 93 + 33);
					mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,ch | COLOR_PAIR(2));
				} else {
					mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,pwchar | COLOR_PAIR(2));
				}
			} else {
				mvaddch(LINES/2,COLS/2-pwfieldlength/2+i,' ' | COLOR_PAIR(2));
			}
		}
		for (int i=0; i<COLS; i++) {
			mvaddch(LINES/2+2,i,' ');
		}
		if (pw.length()<pwfieldlength) {
			move(LINES/2,COLS/2-pwfieldlength/2+pw.length());
		} else {
			move(LINES/2,COLS/2+pwfieldlength/2);
		}
		usleep(WAIT_TICK);
		refresh();
	}
	return "";
}


int main(int argc, char* argv[]) {

	srand(time(0));

	quitmsgs.push_back("There is no escape.");
	quitmsgs.push_back("No quitting.");
	quitmsgs.push_back("Ctrl + C wont save you.");
	quitmsgs.push_back("Try once more?");
	quitmsgs.push_back("Maby it'll work next time.");
	quitmsgs.push_back("Dont even bother.");
	quitmsgs.push_back("No password, no sauce.");
	quitmsgs.push_back("No. Just... No.");
	quitmsgs.push_back("User used QUIT. It is not very effective.");
	quitmsgs.push_back("Gotta try something else bud.");
	quitmsgs.push_back("You need to try harder!");
	quitmsgs.push_back("Pathetic.");

	wrongmsgs.push_back("That password was wrong.");
	wrongmsgs.push_back("Wrong password.");
	wrongmsgs.push_back("Try again. I dare you!");
	wrongmsgs.push_back("Wow that better not be your real password");
	wrongmsgs.push_back("Nope.");
	wrongmsgs.push_back("Try harder!");
	wrongmsgs.push_back("That was almost correct. Not.");
	wrongmsgs.push_back("Please let this not be your actual password.");
	wrongmsgs.push_back("Seriously?");

	signal(SIGINT, handler);
	signal(SIGTSTP, handler);
	signal(SIGABRT, handler);
	signal(SIGSTOP, handler);
	signal(SIGQUIT, handler);

	struct passwd *pw = getpwuid(getuid());
	string homedir=pw->pw_dir;
	pwpath=homedir + "/.xlockpw";

	if (argc<=1) {
		init();
		run();
	} else if (argc==2) {
		if (string(argv[1])=="chpw") {
			changepassword=true;
			fstream f1(pwpath, ios::in);
			if (f1.good()) {
				init();
				run();
			}
			f1.close();
			string p1;
			string p2;
			/*
			cout << endl << "Enter new password: ";
			getline(cin,p1);
			cout << endl << "Repeat password: ";
			getline(cin,p2);
			*/
			p1=readPw("Enter a Password:");
			p2=readPw("Repeat Password:");
			if (p1==p2) {
				ofstream f;
				f.open(pwpath, std::ofstream::out | std::ofstream::trunc);
				if (f.good()) {
					f << sha512(p1);
				} else {
					cerr << "Error while loading password file." << endl;
					cerr << "170" << endl;
					exit(1);
				}
				chmod((char *)pwpath.c_str(),0600);
				f.close();
				cout << endl << "Password updated." << endl;
			} else {
				cerr << "Passwords do not match";
			}
		} else if (string(argv[1])=="soft") {
			softlock=true;
			init();
			run();	
		} /*else if (string(argv[1])=="read") {
			cout << readPw("Enter Password:") << endl;	
		}*/ else {
			cerr << "unknown argument." << endl;
		}
	} else {
		cerr << "too many arguments" << endl;
	}
}