#include <err.h>
#include <errno.h>
#include <fcntl.h>
#include <fnmatch.h>
#include <ftw.h>
#include <limits.h>
#include <sha1.h>
#include <stdarg.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <sys/wait.h>

#define VCX_DIR   ".vcx"
#define HEAD      VCX_DIR "/head"
#define TMP_DIR   VCX_DIR "/tmp"
#define OBJ_DIR   VCX_DIR "/obj"

#define EXCLUDE_PATHS  ".vcxignore" 

#define BUFLEN        8192
#define MAX_DEPTH      256
#define MAX_PATH_LEN  4096

#define KIND_DEL  'D'
#define KIND_MOD  'M'
#define KIND_NEW  'N'

#define MALLOC(s)      _xmalloc((s), __FILE__, __LINE__)
#define REALLOC(p, s)  _xrealloc((p), (s), __FILE__, __LINE__)

static inline void init(int argc, char *argv[]);
static inline void status(int argc, char *argv[]);
static inline void add(int argc, char *argv[]);

static inline void init_stg_area(void);
static inline int cmp_links(const char *link1, const char *link2);
static inline int cmp_files(const char *file1, const char *file2);

static inline void stage_file(const char *file, char kind, int islnk);
static inline void print_status(const char *file, char kind, int islnk);

static inline int scan_head(const char * path, const struct stat *st, int flag, struct FTW *ftwbuf);
static inline int scan_tree(const char * path, const struct stat *st, int flag, struct FTW *ftwbuf);

static inline void mkdirs(const char *path);
static inline int copy_file(const char *src, const char *dst);
static inline void concat(char *dst, size_t n, const char *arg1, ...);

static inline void *_xmalloc(size_t s, const char *file, int line);
static inline void *_xrealloc(void *ptr, size_t s, const char *file, int line);

struct command {
	char *name;
	void (*func)(int argc, char *argv[]);
};

struct command cmd[] = {
	{"init", init}, 
	{"add", add}, 
	{"status", status}, 
	{NULL, NULL}
};

int main(int argc, char *argv[])
{
	if (argc < 2)
		errx(1, "Usage: %s <command> [<args>]", argv[0]);

	for (int i = 0; cmd[i].name != NULL; i++) {
		if (strcmp(argv[1], cmd[i].name) == 0) {
			cmd[i].func(argc - 1, argv + 1);
			return 0;
		}
	}
	return 0;
}

static inline void init(int argc, char *argv[])
{
	if (mkdir(VCX_DIR, 0755) == -1) {
		if (errno != EEXIST)
			err(1, "Failed to create %s", VCX_DIR);
	}

	if (mkdir(HEAD, 0755) == -1) {
		if (errno != EEXIST)
			err(1, "Failed to create %s", HEAD);
	}

	if (mkdir(OBJ_DIR, 0755) == -1) {
		if (errno != EEXIST)
			err(1, "Failed to create %s", OBJ_DIR);
	}

	printf("Repository ready\n");
}

static void (*scan_res_cb)(const char *file, char kind, int islnk) = NULL;

static inline void status(int argc, char *argv[])
{
	scan_res_cb = print_status;

	if (nftw(HEAD, scan_head, MAX_DEPTH, FTW_PHYS) == -1)
		err(1, "Failed to scan head");	

	if (nftw(".", scan_tree, MAX_DEPTH, FTW_PHYS) == -1)
		err(1, "Failed to scan work tree");	
}

static inline void add(int argc, char *argv[])
{
	size_t i;
	struct stat st;
	char *wt_path;
	
	if (argc < 2)
		errx(1, "Usage: %s [<files>]", argv[0]);

	init_stg_area();
	scan_res_cb = stage_file;

	for (i = 1; i < argc; i++) {
		wt_path = argv[i];
		if (lstat(wt_path, &st) == -1)
			err(1, "%s", wt_path);

		size_t hd_sz = strlen(HEAD) + strlen(wt_path) + 2;
		char hd_path[hd_sz];
		concat(hd_path, hd_sz, HEAD, "/", wt_path, NULL);

		size_t tmp_sz = strlen(TMP_DIR) + strlen(wt_path) + 2;
		char tmp_path[tmp_sz];
		concat(tmp_path, tmp_sz, TMP_DIR, "/", wt_path, NULL);
		
		if (nftw(hd_path, scan_head, MAX_DEPTH, FTW_PHYS) == -1)
			err(1, "Failed to scan head");	

		if (nftw(wt_path, scan_tree, MAX_DEPTH, FTW_PHYS) == -1)
			err(1, "Failed to scan work tree");	
	}
}

int scan_head(const char * path, const struct stat *sb, int flag, 
	struct FTW *ftwbuf)
{
	struct stat wt_stat;

	if (flag == FTW_F) {
		const char *filename = path + ftwbuf->base;

		size_t wt_sz = strlen(filename) + 3;
		char wt_path[wt_sz];
		concat(wt_path, wt_sz, ".", "/", filename, NULL);

		if (lstat(wt_path, &wt_stat) == -1) {
			if (errno == ENOENT) {
				// file/link doesn't matter for deletions
				scan_res_cb(wt_path, KIND_DEL, 0);
				return 0;
			}
			else
				err(1, "lstat failed for %s", wt_path);
		}

		if (wt_stat.st_size == sb->st_size && 
			wt_stat.st_mtime == sb->st_mtime)
			return 0; // no change

		if (S_ISREG(wt_stat.st_mode)) {
			if (cmp_files(path, wt_path) == 1)
				scan_res_cb(wt_path, KIND_MOD, 0);
		} else if (S_ISLNK(wt_stat.st_mode)) {
			if (cmp_links(path, wt_path) == 1)
				scan_res_cb(wt_path, KIND_MOD, 1);
		}
	}

	return 0;
}

static inline int scan_tree(const char * path, const struct stat *st, 
	int flag, struct FTW *ftwbuf)
{
	if (fnmatch("./.vcx/*", path, 0) == 0)
		return 0;

	if (flag == FTW_F) {
		const char *filename = path + ftwbuf->base;

		size_t hd_sz = strlen(HEAD) + strlen(filename) + 2;
		char hd_path[hd_sz];
		concat(hd_path, hd_sz, HEAD, "/", filename, NULL);

		if (access(hd_path, F_OK) != 0)
			scan_res_cb(path, KIND_NEW, S_ISLNK(st->st_mode) != 0);
	}

	return 0;
}

static inline int cmp_files(const char *file1, const char *file2)
{
	int rc;
	int fd1, fd2;
	ssize_t r1, r2;
	char buf1[BUFLEN], buf2[BUFLEN];

	if ((fd1 = open(file1, O_RDONLY)) < 0)
		err(1, "Couldn't open %s to compare", file1);

	if ((fd2 = open(file2, O_RDONLY)) < 0)
		err(1, "Couldn't open %s to compare", file2);

	rc = 0;
	do {
		r1 = read(fd1, buf1, BUFLEN);
		r2 = read(fd2, buf2, BUFLEN);
		
		if (r1 != r2 || memcmp(buf1, buf2, r1) != 0) {
			rc = 1;
			break;
		} 
	} while (r1 > 0);

	close(fd1);
	close(fd2);
	return rc;
}

static inline int cmp_links(const char *link1, const char *link2)
{
	ssize_t len1, len2; 
	char buf1[PATH_MAX], buf2[PATH_MAX];	
	
	len1 = readlink(link1, buf1, sizeof(buf1) - 1);
	len2 = readlink(link2, buf2, sizeof(buf2) - 1);
	
	if (len1 < 0 || len2 < 0)
		err(1, "readlink error on %s or %s", link1, link2);

	if (len1 != len2)
		return 1;

	buf1[len1] = '\0';
	buf2[len2] = '\0';

	return strcmp(buf1, buf2);
}

static inline void init_stg_area(void)
{
	pid_t rm_pid, cp_pid;

	rm_pid = fork();
	if (rm_pid == 0) {
		if (access(TMP_DIR, F_OK) == 0) {
			execlp("rm", "rm", "-rf", TMP_DIR, NULL);
			err(1, "Failed to delete %s", TMP_DIR);
		}
		_exit(0);
	} else {
		waitpid(rm_pid, NULL, 0);
		cp_pid = fork();
		if (cp_pid == 0) {
			execlp("cp", "cp", "-RP", HEAD, TMP_DIR, NULL);
			err(1, "Failed to copy %s to %s", HEAD, TMP_DIR);
		}
		waitpid(cp_pid, NULL, 0);
	}
}

static inline void stage_file(const char *file, char kind, int islnk)
{
	int len, obj_len;
	char obj_name[SHA1_DIGEST_STRING_LENGTH];

	switch (kind) {
	case KIND_MOD:
		break;
	case KIND_NEW:
		if (!islnk) {
			len = strlen(file);
			obj_len = SHA1_DIGEST_STRING_LENGTH - 1;
			
			if (SHA1Data((const uint8_t *)file, len, obj_name) == NULL)
				err(1, "Failed to compute hash for %s", file);

			size_t obj_sz = strlen(OBJ_DIR) + obj_len + 6;
			char obj_path[obj_sz];
			concat(obj_path, obj_sz, OBJ_DIR, "/", obj_name, ".tmp", NULL);
			copy_file(file, obj_path);

			// create link relative to tmp/
			size_t lnkt_sz = strlen("../obj/") + obj_len + 5;
			char lnk_target[lnkt_sz];
			concat(lnk_target, lnkt_sz, "../obj/", obj_name, ".tmp", NULL);

			size_t lnkp_sz = strlen(TMP_DIR) + len + 2;
			char lnk_path[lnkp_sz];
			concat(lnk_path, lnkp_sz, TMP_DIR, "/", file, NULL);

			mkdirs(lnk_path);
			if (symlink(lnk_target, lnk_path) != 0)
				err(1, "symlink error on %s", file);
		} else {
			
		}
		break;
	case KIND_DEL:
		break;
	default:
		break;
	}
}

static inline void print_status(const char *file, char kind, int islnk)
{
	printf("[%c] %s\n", kind, file);
}

static inline void concat(char *dst, size_t n, const char *arg1, ...)
{
	va_list ap;
	const char *part;
	int curlen, partlen;

	if (!dst || n == 0)
		return;

	dst[0] = '\0';
	curlen = 0;

	if (!arg1)
		return;

	va_start(ap, arg1);
	part = arg1;

	while (part) {
		partlen = strlen(part);
		if (curlen + partlen + 1 > n) {
			dst[curlen] = '\0'; 
			va_end(ap);
			errx(1, "Path too long: %s%s", dst, part);
		}
		memcpy(dst + curlen, part, partlen);
		curlen += partlen;
		part = va_arg(ap, const char *);
	}

	dst[curlen] = '\0';
	va_end(ap);
}

static inline int diff(const char *file1, const char *file2)
{
	pid_t pid;
	int null_fd, status, rc;

	pid = fork();

	if (pid == 0) {
		null_fd = open("/dev/null", O_WRONLY);
		if (null_fd == -1)
			err(1, "open() failed on /dev/null");

		dup2(null_fd, STDOUT_FILENO);
		dup2(null_fd, STDERR_FILENO);
		close(null_fd);

		execlp("diff", "diff", "-q", file1, file2, (char *)NULL);
		err(1, "execlp failed for diff");
	} else if (pid > 0) {
		waitpid(pid, &status, 0);
		// diff returns 0 if files are same, 1 if different, >1 if error
		rc = WEXITSTATUS(status);
		if (rc > 1)
			err(1, "diff error");
		return rc;
	}

	err(1, "fork");
}

static inline int copy_file(const char *src, const char *dst)
{
	int fdin, fdout;
	struct stat st;
	char buf[BUFLEN], *outptr;
	ssize_t nread, nwrite, res;

	if ((fdin = open(src, O_RDONLY)) < 0)
		err(1, "open failed for %s", src);

	if (fstat(fdin, &st) < 0) {
		close(fdin);
		err(1, "fstat failed for %s", src);
	}

	mkdirs(dst);

	fdout = open(dst, O_WRONLY | O_CREAT | O_TRUNC, st.st_mode & 0777);
	if (fdout < 0) {
		close(fdin);
		err(1, "open failed for %s", dst);
	}

	while ((nread = read(fdin, buf, sizeof(buf))) > 0) {
		nwrite = 0;
		outptr = buf;
		
		// Handle partial writes (possible with large files/slow disks)
		while (nwrite < nread) {
			res = write(fdout, outptr + nwrite, nread - nwrite);
			if (res < 0) {
				if (errno == EINTR) // interrupted by a signal - retry
					continue;
				close(fdin);
				close(fdout);
				err(1, "Copy error on %s", src); // bad error: e.g, disk full
			}
			nwrite += res;
		}
	}

	close(fdin);
	close(fdout);
	return (nread < 0) ? -1 : 0;
}

static inline void mkdirs(const char *path)
{
	char *p;
	int pathlen;

	static char buf[MAX_PATH_LEN];

	pathlen = strlen(path);
	if (MAX_PATH_LEN < pathlen + 1)
		errx(1, "Path too long: %s", path);

	buf[0] = '\0';
	strcpy(buf, path);

	for (p = buf; *p; p++) {
		if (*p == '/') {
			*p = '\0';
			if (mkdir(buf, 0755) != 0) {
				if (errno != EEXIST)
					err(1, "mkdir failed for %s", buf);
			}
			*p = '/';
		}
	}
}

static inline void *_xmalloc(size_t s, const char *file, int line)
{
	void *p;

	if (!(p = malloc(s)))
		err(1, "%s:%d: malloc", file, line);
	return p;
}

static inline void *_xrealloc(void *ptr, size_t s, const char *file, 
	int line)
{
	void *p;

	if (!(p = realloc(ptr, s)))
		err(1, "%s:%d: realloc", file, line);
	return p;
}