#include <windows.h>
#include <stdlib.h>
#include <tchar.h>
#include <winapi/winbase.h> /* Contains GetFileSizeEx */
#ifndef UNICODE
# include <stdio.h>
#else
# include <wchar.h>
#endif
#include "sqlite3.h"
#include "sha-256.h"
#define SHASIZ SIZE_OF_SHA_256_HASH

/* TODO: "\\?\" trick?
 * Windows limits path lengths to 260 characters. Explorer on Windows XP
 * has this issue too.
 *
 * UTF issues?
 */

/*************************************************************************
*  Compatability defines for non-Unicode systems
 ************************************************************************/

#ifdef UNICODE
# define ENCODING SQLITE3_UTF16
# define sqlite3_open_U sqlite3_open16
# define sqlite3_errmsg_U sqlite3_errmsg16
# define sqlite3_bind_text_U sqlite3_bind_text16
# define sqlite3_column_bytes_U sqlite3_column_bytes16
# define sqlite3_column_text_U sqlite3_column_text16
#else
# define ENCODING SQLITE3_UTF8
# define sqlite3_open_U sqlite3_open
# define sqlite3_errmsg_U sqlite3_errmsg
# define sqlite3_bind_text_U sqlite3_bind_text
# define sqlite3_column_bytes_U sqlite3_column_bytes
# define sqlite3_column_text_U sqlite3_column_text
#endif

/****************
 * Globals
 ***************/

sqlite3 *g_db = NULL;
enum {
	NORMAL_MODE,
	VERBOSE_MODE,
	DEBUG_MODE
} g_verbose = NORMAL_MODE;
sqlite3_int64 g_backup_id = -1;

/********************
 * Utility Functions
 *******************/

static void _log(TCHAR *msg, int verb, ...)
{
	va_list va;

	if (g_verbose >= verb) {
		va_start(va, verb);
		_vftprintf(stderr, msg, va);
		va_end(va);
	}
}
/* String literals only */
#define log(msg, ...) _log(TEXT(msg), __VA_ARGS__)

static void _die(TCHAR *emsg, ...)
{
	va_list va;
	va_start(va, emsg);
	_vtprintf(emsg, va);
	va_end(va);

	sqlite3_close(g_db);
	exit(1);
}
/* String literals only */
#define die(emsg, ...) _die(TEXT(emsg), __VA_ARGS__)

/**********************************
 * Database constants
 *********************************/

/* 'wlb' in ASCII */
#define UPPER_VERSION 0x776c62
#define CUR_VERSION 0x776c6201
#define S2(s) #s
#define S(s) S2(s)
static const char *init_script =
"BEGIN;\n"
"PRAGMA user_version = " S(CUR_VERSION) ";\n"
"CREATE TABLE chunks (\n"
"	rowid INTEGER PRIMARY KEY, \n"
"	sha256 TEXT UNIQUE NOT NULL, \n"
"	data BLOB NOT NULL\n"
"); \n"
"CREATE TABLE backup_ids (\n"
"	ts TEXT UNIQUE NOT NULL,\n"
"	rowid INTEGER PRIMARY KEY\n"
"); \n"
"CREATE TABLE backups (\n"
"	backup INTEGER NOT NULL REFERENCES backup_ids ON DELETE CASCADE, \n"
"	path TEXT NOT NULL, \n"
"	chunk INTEGER NOT NULL REFERENCES chunks ON DELETE RESTRICT\n"
"); \n"
"CREATE INDEX backups_paths ON backups (path); \n"
"CREATE INDEX backups_backup ON backups (backup); \n"
"CREATE INDEX backups_chunk ON backups (chunk);\n"
"COMMIT;";

/* Check if the database version is correct. If it is not correct
 * and the create flag is enabled, then initialize a new database.
 */
static void check_db_version(int did_not_exist)
{
	int user_version;
	sqlite3_stmt *stmt;

	/* DB did not exist beforehand: initialize DB */
	if (did_not_exist) {
		log("Writing init script %s\n", DEBUG_MODE, init_script);
		if (sqlite3_exec(g_db, init_script, NULL, NULL, NULL) != SQLITE_OK)
			die("Error initializing database: %s\n",
			    sqlite3_errmsg_U(g_db));
		log("%s", VERBOSE_MODE, "Intialized database\n");
		return;
	}

	/* DB existed beforehand: do version check */
	if (sqlite3_prepare_v2(g_db, "PRAGMA user_version;", -1, &stmt, NULL)
	                       != SQLITE_OK)
		die("Error preparing user version check: %s\n",
		    sqlite3_errmsg_U(g_db));
	if (sqlite3_step(stmt) != SQLITE_ROW)
		die("Error executing user version check: %s\n",
		    sqlite3_errmsg_U(g_db));
	user_version = sqlite3_column_int(stmt, 0);
	sqlite3_finalize(stmt);

	if (user_version != CUR_VERSION) {
		die("Bad DB version found (expected %d), got %d\n",
		    CUR_VERSION & 0xFF, user_version & 0xFF);
	}
}

/* Initialize the backup ID used in this session.
 * This is only called if the user requests an archive of a directory.
 */
static void initialize_backup_id(void)
{
	sqlite3_stmt *stmt;

	if (sqlite3_prepare(g_db,
	                    "INSERT INTO backup_ids (ts) VALUES (datetime('subsec')) RETURNING ts,rowid;",
	                    -1, &stmt, NULL) != SQLITE_OK) {
		die("failed to prepare inserting timestamp: %s\n",
		    sqlite3_errmsg_U(g_db));
	}

	if (sqlite3_step(stmt) != SQLITE_ROW)
		die("failed to insert timestamp: %s\n",
		    sqlite3_errmsg_U(g_db));

	log("Timestamp: %s\n", NORMAL_MODE, sqlite3_column_text_U(stmt, 0));
	g_backup_id = sqlite3_column_int64(stmt, 1);

	sqlite3_finalize(stmt);
}

/* Opens the DB and initializes it if it did not exist. */
static void open_db(TCHAR *fn)
{
	int did_not_exist = GetFileAttributes(fn)
	                    == INVALID_FILE_ATTRIBUTES;
	char *errmsg;

	log("Opening DB %s\n", DEBUG_MODE, fn);
	if (sqlite3_open_U(fn, &g_db) != SQLITE_OK)
		die("Error opening database %s: %s\n", fn,
		    sqlite3_errmsg_U(g_db));
	check_db_version(did_not_exist);

	if (sqlite3_exec(g_db, "BEGIN;", NULL, NULL, &errmsg) != SQLITE_OK)
		die("failed to begin transaction: %s\n", errmsg);
	log("Started transaction\n", DEBUG_MODE);
}

/* Insert a backup record into the database.
 *
 * This part occurs after the file data has been inserted or identified
 * by SHA256 hash.
 */
static int insert_backup_record(TCHAR *name, sqlite_int64 chunk_id)
{
	sqlite3_stmt *stmt;
	int r = 0;

	log("Inserting backup record (%s, %lld)\n", DEBUG_MODE, name, (long long) chunk_id);

	if (sqlite3_prepare(g_db,
	                    "INSERT INTO backups (backup, path, chunk) VALUES (?,?,?);",
	                    -1, &stmt, NULL) != SQLITE_OK) {
		log("Could not prepare backup insert statement for %s: %s\n",
		    NORMAL_MODE, name, sqlite3_errmsg_U(g_db));
		return 0;
	}

	if (sqlite3_bind_int64(stmt, 1, g_backup_id) != SQLITE_OK) {
		log("Could not bind backup id for %s: %s\n", NORMAL_MODE,
		     name, sqlite3_errmsg_U(g_db));
		goto end;
	}

	if (sqlite3_bind_text_U(stmt, 2, name, -1, SQLITE_TRANSIENT) != SQLITE_OK) {
		log("Could not bind path for %s: %s\n", NORMAL_MODE,
		    name, sqlite3_errmsg_U(g_db));
		goto end;
	}

	if (sqlite3_bind_int64(stmt, 3, chunk_id) != SQLITE_OK) {
		log("Could not bind chunk for %s: %s\n", NORMAL_MODE, name,
		         sqlite3_errmsg_U(g_db));
		goto end;
	}

	if (sqlite3_step(stmt) != SQLITE_DONE) {
		log("Error in inserting chunk for %s: %s\n", NORMAL_MODE, name,
		         sqlite3_errmsg_U(g_db));
		goto end;
	}

	r = 1;
	log("Sucessful insertion\n", DEBUG_MODE);
end:
	sqlite3_finalize(stmt);
	return r;
}

/* Write a file to a BLOB in ``chunks`` at row ``rowid``.
 *
 * Returns 0 on failure and 1 on success.
 */
static int write_chunk(TCHAR *name, HANDLE f, sqlite_int64 rowid)
{
	sqlite3_blob *blob;
	char buf[4096];
	DWORD read = 0;
	int has_written = 0;
	int r = 0;

	log("Writing chunk for %s, rowid=%lld\n", DEBUG_MODE, name, (long long)rowid);

	if (sqlite3_blob_open(g_db, "main", "chunks", "data", rowid, 1,
	                      &blob) != SQLITE_OK) {
		log("Could not open BLOB to write chunk for %s: %s\n", NORMAL_MODE,
		         name, sqlite3_errmsg_U(g_db));
		return 0;
	}

	for (;;) {
		if (ReadFile(f, buf, sizeof(buf), &read, NULL) == FALSE) {
			log("failed to read in %s for sha256 calculation\n", NORMAL_MODE,
			         name);
			goto end;
		}

		if (read == 0)
			break;
		if (sqlite3_blob_write(blob, buf, read, has_written)
		    != SQLITE_OK) {
			log("Could not write to BLOB for %s: %s\n", NORMAL_MODE,
			         name, sqlite3_errmsg_U(g_db));
			goto end;
		}

		has_written += read;
		log("Blob written %d bytes\n", DEBUG_MODE, has_written);
	}

	r = 1;
	log("Succesfully written blob\n", DEBUG_MODE);
end:
	sqlite3_blob_close(blob);
	return r;
}

/* Insert a chunk with a specified SHA256 checksum into the database.
 *
 * Returns 1 on success, 0 on failure. On success, ``rowid`` contains
 * the row in ``chunks`` that contains the data of ``f``.
 */
static int insert_chunk(TCHAR *name, HANDLE f, sqlite_int64 *rowid,
                        uint8_t sha256[SHASIZ])
{
	/* NOTE: Windows XP x86 does not support GetFilesizeEx() */
	DWORD fsizeLow, fsizeHigh;
	sqlite3_int64 fsize;
	sqlite3_stmt *stmt;
	int r = 0;

	if (SetFilePointer(f, 0, NULL, FILE_BEGIN)
	    == INVALID_SET_FILE_POINTER) {
		log("rewind for %s failed\n", NORMAL_MODE, name);
		return 0;
	}

	fsizeLow = GetFileSize(f, &fsizeHigh);
	if (fsizeLow == INVALID_FILE_SIZE) {
		log("Could not get the file size of %s\n", NORMAL_MODE, name);
		return 0;
	}
	fsize = fsizeLow | (fsizeHigh << 32);
	log("File size: %lld\n", (long long)fsize);

	/* Chunks are fixed size in SQLite. They need to be pre-allocated
	 * with the size of a file before a file can be written to it.
	 */
	if (sqlite3_prepare(g_db,
	                    "INSERT INTO chunks (sha256, data) VALUES (hex(?), zeroblob(?)) RETURNING rowid;",
	                    -1, &stmt, NULL) != SQLITE_OK) {
		log("Could not prepare chunk insertion for %s: %s\n", NORMAL_MODE,
		         name, sqlite3_errmsg_U(g_db));
		return 0;
	}

	if (sqlite3_bind_blob(stmt, 1, sha256, SHASIZ, SQLITE_TRANSIENT)
	    != SQLITE_OK) {
		log("Could not bind sha256 to statement for %s: %s\n", NORMAL_MODE,
		         name, sqlite3_errmsg_U(g_db));
		goto finalize;
	}

	if (sqlite3_bind_int64(stmt, 2, fsize) != SQLITE_OK) {
		log("Could not bind file size to statement for %s: %s\n", NORMAL_MODE,
		         name, sqlite3_errmsg_U(g_db));
		goto finalize;
	}

	if (sqlite3_step(stmt) != SQLITE_ROW) {
		log("Could not step chunk insertion statement for %s: %s\n", NORMAL_MODE,
		         name, sqlite3_errmsg_U(g_db));
		goto finalize;
	}

	*rowid = sqlite3_column_int64(stmt, 0);
	log("Allocated new chunk at %lld\n", DEBUG_MODE, (long long)*rowid);
	r = write_chunk(name, f, *rowid);
finalize:
	sqlite3_finalize(stmt);
	return r;
}

/* Print SHA256 sum ASCII representation to terminal (debug mode only) */
static void dump_sha256(uint8_t sha256[SHASIZ])
{
	int i;

	for (i = 0; i < SHASIZ; i++) {
		log("%hhX", DEBUG_MODE, sha256[i]);
	}
}

/* Check if the SHA256 checksum already exists in the database.
 *
 * If the checksum exists, ``rowid`` is filled with the row in ``chunks``
 * that contains that sha256 sum. If it does not exist, then ``rowid``
 * contains ``-1``.
 *
 * Returns 0 if an error occured, and 1 if no error occured.
 */
static int check_sha256(TCHAR *name, uint8_t sha256[SHASIZ],
                        sqlite_int64 *rowid)
{
	sqlite3_stmt *stmt;
	int r = 0;

	/* Since ``sha256`` is binary, use ``hex`` to convert the binary
	 * string to a text representation.
	 */
	if (sqlite3_prepare_v2(g_db,
	                       "SELECT rowid FROM chunks WHERE sha256 = hex(?);",
	                       -1, &stmt, NULL) != SQLITE_OK) {
		log("failed to prepare sha256 statement for %s: %s\n", NORMAL_MODE,
		         name, sqlite3_errmsg_U(g_db));
		goto end;
	}

	if (sqlite3_bind_blob(stmt, 1, sha256, SHASIZ, SQLITE_TRANSIENT)
	    != SQLITE_OK) {
		log("failed to bind sha256 value for %s: %s\n", NORMAL_MODE,
		         name, sqlite3_errmsg_U(g_db));
		goto end;
	}

	switch (sqlite3_step(stmt)) {
	case SQLITE_ROW:
		/* there is a sha256 value */
		*rowid = sqlite3_column_int(stmt, 0);
		break;
	case SQLITE_DONE:
		/* The chunk was never entered */
		*rowid = -1;
		break;
	}

	r = 1;
end:
	log("rowid: %lld\n", DEBUG_MODE, (long long)*rowid);
	sqlite3_finalize(stmt);
	return r;
}

/* Calculate the SHA256 checksum of the file in ``f``, and place a binary
 * representation of the checksum into ``shastr``.
 *
 * Returns 0 if there was an error, 1 for success.
 */
static int calculate_sha256(TCHAR *name, HANDLE f, uint8_t sha256[SHASIZ])
{
	DWORD read = 0;
	uint8_t buf[SIZE_OF_SHA_256_CHUNK * 64];
	struct Sha_256 sha_state;
	int i;

	sha_256_init(&sha_state, sha256);
	for (;;) {
		if (ReadFile(f, buf, sizeof(buf), &read, NULL) == FALSE) {
			log("failed to read in %s for sha256 calculation\n", NORMAL_MODE,
			         name);
			return 0;
		}

		/* ReadFile() reads 0 bytes on EOF. */
		if (read == 0)
			break;
		sha_256_write(&sha_state, buf, read);
	}

	log("SHA256 sum of %s: ", DEBUG_MODE, name);
	dump_sha256(sha256);
	log("\n", DEBUG_MODE);
	sha_256_close(&sha_state);
	return 1;
}

/* Archive a file. */
static void archive_file(TCHAR *name)
{
	HANDLE f;
	sqlite_int64 chunks_rowid = -1;
	uint8_t sha256[SHASIZ];
	char *errstr;
	int success = 0;

	if (sqlite3_exec(g_db, "SAVEPOINT file;", NULL, NULL, &errstr)
	    != SQLITE_OK)
		die("failed to initialize file savepoint: %s\n", errstr);

	f = CreateFile(name, GENERIC_READ, FILE_SHARE_READ, NULL,
	                     OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, NULL);
	if (f == INVALID_HANDLE_VALUE) {
		_tprintf("could not open %s\n", name);
		goto end;
	}

	if (!calculate_sha256(name, f, sha256))
		goto end;
	if (!check_sha256(name, sha256, &chunks_rowid))
		goto end;

	if (chunks_rowid < 0) {
		if (!insert_chunk(name, f, &chunks_rowid, sha256))
			goto end;
	} else {
		log("SHA sum already exists\n", DEBUG_MODE);
	}

	if (!insert_backup_record(name, chunks_rowid))
		goto end;
	success = 1;

	log("Archived %s\n", VERBOSE_MODE, name);
end:
	if (sqlite3_exec(g_db, success ? "RELEASE file;" : "ROLLBACK TO file;", NULL, NULL,
	                 &errstr) != SQLITE_OK)
		die("failed to rollback file savepoint: %s\n", errstr);
	CloseHandle(f);
}

/* Recursively archive a directory. */
static void archive_directory(TCHAR *dirname)
{
	WIN32_FIND_DATA fdata;
	HANDLE dhandle;
	TCHAR pathname[PATH_MAX];

	if (_sntprintf(pathname, PATH_MAX, TEXT("%s\\*"), dirname) < 0) {
		log("Pathname %s\\* too long\n", NORMAL_MODE, dirname);
		return;
	}

	dhandle = FindFirstFile(pathname, &fdata);
	if (dhandle == INVALID_HANDLE_VALUE) {
		log("Failed to open directory %s\n", NORMAL_MODE, dirname);
		return;
	}
	log("Archiving directory %s\n", VERBOSE_MODE, dirname);

	do {
		/* Windows versions prior to 11 (!) include _snprintf() as
		 * a non-standard version of snprintf(). _snprintf() will
		 * return -1 if the string does not fit in the buffer.
		 */
		if (_sntprintf(pathname, PATH_MAX, TEXT("%s\\%s"), dirname,
		               fdata.cFileName) < 0) {
			log("Pathname for %s\\%s too long\n", NORMAL_MODE, dirname,
			         fdata.cFileName);
			continue;
		}

		if (_tcscmp(fdata.cFileName, TEXT(".")) == 0
		    || _tcscmp(fdata.cFileName, TEXT("..")) == 0) {
			log("Skipping . or ..\n", DEBUG_MODE);
			continue;
		}

		if (fdata.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY)
			archive_directory(pathname);
		else
			archive_file(pathname);
	} while (FindNextFile(dhandle, &fdata) != 0);

	FindClose(dhandle);
}

/* Check if the path is a file or a directory, and either archive
 * the file or archive the entire directory recursively. */
static void archive_file_or_directory(TCHAR *name)
{
	DWORD attr = GetFileAttributes(name);
	if (attr == INVALID_FILE_ATTRIBUTES) {
		log("failed to open %s\n", NORMAL_MODE, name);
		return;
	}

	if (g_backup_id == -1)
		initialize_backup_id();

	if (attr & FILE_ATTRIBUTE_DIRECTORY)
		archive_directory(name);
	else
		archive_file(name);
}

/* Print instructions to console and exit. */
static void usage(void)
{
	_tprintf(TEXT("\nwlb [\\V] [\\VV] [\\S] [\\H] [\\D DBNAME] [\\A DIRECTORIES...]\n"));
	_tprintf(TEXT("\\V: Verbose\n"));
	_tprintf(TEXT("\\VV: Debug mode\n"));
	_tprintf(TEXT("\\S: Go back to normal mode\n"));
	_tprintf(TEXT("\\H: Display the help\n"));
	_tprintf(TEXT("\\D: Database (create if does not exist)\n"));
	_tprintf(TEXT("\\A: Archive directories\n"));
	exit(0);
}

int _tmain(int argc, TCHAR *argv[])
{
	int i;

	for (i = 1; i < argc; i++) {
		if (_tcscmp(argv[i], TEXT("\\H")) == 0) {
			usage();
		} else if (_tcscmp(argv[i], TEXT("\\A")) == 0) {
			i++;
			archive_file_or_directory(argv[i]);
		} else if (_tcscmp(argv[i], TEXT("\\D")) == 0) {
			i++;
			open_db(argv[i]);
		} else if (_tcscmp(argv[i], TEXT("\\V")) == 0) {
			g_verbose = VERBOSE_MODE;
		} else if (_tcscmp(argv[i], TEXT("\\VV")) == 0) {
			g_verbose = DEBUG_MODE;
		} else if (_tcscmp(argv[i], TEXT("\\S")) == 0) {
			g_verbose = NORMAL_MODE;
		} else {
			usage();
		}
	}

	if (g_db)
		sqlite3_exec(g_db, "COMMIT;", NULL, NULL, NULL);
	sqlite3_close(g_db);
	return 0;
}