#include "gc.h"
#include <errno.h>
#include <fcntl.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/stat.h>
#include <unistd.h>

struct publication {
	const char *dst;
	char *stage;
	char *backup;
	int had_old;
	int installed;
	int passthrough;
};

static char *
publish_path(const char *dst, const char *kind)
{
	int n = snprintf(NULL, 0, "%s.w6c.%ld.%s", dst, (long)getpid(), kind);
	if (n < 0 || (size_t)n == (size_t)-1) return NULL;
	char *p = malloc((size_t)n + 1);
	if (p == NULL) return NULL;
	if (snprintf(p, (size_t)n + 1, "%s.w6c.%ld.%s", dst,
	    (long)getpid(), kind) != n) {
		free(p);
		return NULL;
	}
	return p;
}

static int
copy_stream(FILE *src, FILE *out)
{
	int bad = fflush(src) != 0 || fseek(src, 0, SEEK_SET) != 0;
	unsigned char buf[65536];
	while (!bad) {
		size_t n = fread(buf, 1, sizeof buf, src);
		if (n != 0 && fwrite(buf, 1, n, out) != n) bad = 1;
		if (n < sizeof buf) {
			if (ferror(src)) bad = 1;
			break;
		}
	}
	return bad ? -1 : 0;
}

static int
materialize_stream(FILE *src, struct publication *p)
{
	struct stat st;
	if (lstat(p->dst, &st) == 0 && S_ISCHR(st.st_mode)) {
		FILE *out = fopen(p->dst, "wb");
		int bad = out == NULL || copy_stream(src, out) < 0;
		if (out != NULL && fclose(out) != 0) bad = 1;
		if (bad) {
			fprintf(stderr, "w6c: cannot write %s\n", p->dst);
			return -1;
		}
		p->passthrough = 1;
		return 0;
	}
	p->stage = publish_path(p->dst, "new");
	p->backup = publish_path(p->dst, "old");
	if (p->stage == NULL || p->backup == NULL) {
		fputs("w6c: out of memory\n", stderr);
		return -1;
	}
	int fd = open(p->stage, O_WRONLY | O_CREAT | O_EXCL, 0644);
	if (fd < 0) {
		fprintf(stderr, "w6c: cannot stage %s\n", p->dst);
		return -1;
	}
	FILE *out = fdopen(fd, "wb");
	if (out == NULL) {
		close(fd);
		unlink(p->stage);
		return -1;
	}
	int bad = copy_stream(src, out) < 0;
	if (fclose(out) != 0) bad = 1;
	if (bad) {
		unlink(p->stage);
		fprintf(stderr, "w6c: cannot stage %s\n", p->dst);
		return -1;
	}
	return 0;
}

static void
publication_free(struct publication *p)
{
	free(p->backup);
	free(p->stage);
}

/* A publication rollback name is occupied by any terminal directory entry,
 * including a dangling symlink.  Never follow it while deciding whether the
 * compiler may park an existing destination there. */
static int
path_exists_nofollow(const char *path)
{
	struct stat st;
	if (lstat(path, &st) == 0) return 1;
	return errno == ENOENT ? 0 : -1;
}

static int
path_is_regular_nofollow(const char *path)
{
	struct stat st;
	return lstat(path, &st) == 0 && S_ISREG(st.st_mode);
}

/* Publish assembly and interface as one rollback group. Named output bytes
 * never change before checking and lowering have both succeeded. */
static int
publish_all(struct publication *p, int n)
{
	for (int i = 0; i < n; i++) {
		if (p[i].passthrough) continue;
		if (!path_is_regular_nofollow(p[i].stage)) {
			fprintf(stderr, "w6c: publication stage is not a regular file for %s\n",
			    p[i].dst);
			goto rollback;
		}
		struct stat st;
		if (lstat(p[i].dst, &st) == 0) {
			if (!S_ISREG(st.st_mode)) {
				fprintf(stderr,
				    "w6c: publication destination is not a regular file: %s\n",
				    p[i].dst);
				goto rollback;
			}
		} else if (errno != ENOENT) {
			fprintf(stderr, "w6c: cannot inspect publication destination %s\n",
			    p[i].dst);
			goto rollback;
		}
	}
	for (int i = 0; i < n; i++) {
		if (p[i].passthrough) continue;
		int exists = path_exists_nofollow(p[i].backup);
		if (exists != 0) {
			if (exists < 0)
				fprintf(stderr, "w6c: cannot inspect publication backup for %s\n",
				    p[i].dst);
			else
			fprintf(stderr, "w6c: publication backup exists for %s\n",
			    p[i].dst);
			goto rollback;
		}
	}
	for (int i = 0; i < n; i++) {
		if (p[i].passthrough) continue;
		if (rename(p[i].dst, p[i].backup) == 0)
			p[i].had_old = 1;
		else if (errno != ENOENT) {
			fprintf(stderr, "w6c: cannot preserve %s\n", p[i].dst);
			goto rollback;
		}
	}
	for (int i = 0; i < n; i++) {
		if (p[i].passthrough) continue;
		if (rename(p[i].stage, p[i].dst) != 0) {
			fprintf(stderr, "w6c: cannot publish %s\n", p[i].dst);
			goto rollback;
		}
		p[i].installed = 1;
	}
	/* Installation is the commit point.  Backup cleanup cannot truthfully
	 * turn a completely installed pair into a rejected compilation. */
	for (int i = 0; i < n; i++)
		if (!p[i].passthrough && p[i].had_old
		    && unlink(p[i].backup) != 0) {
			fprintf(stderr, "w6c: cannot remove backup for %s\n", p[i].dst);
		}
	return 0;

rollback:
	for (int i = n - 1; i >= 0; i--) {
		if (p[i].passthrough) continue;
		if (p[i].installed) (void)unlink(p[i].dst);
		if (p[i].had_old) (void)rename(p[i].backup, p[i].dst);
		if (p[i].stage != NULL) (void)unlink(p[i].stage);
	}
	return -1;
}

static int
slurp(const char *path, char **outbuf, u64 *outlen)
{
	FILE *f = fopen(path, "rb");
	if (f == NULL) return -1;
	fseek(f, 0, SEEK_END);
	long n = ftell(f);
	fseek(f, 0, SEEK_SET);
	if (n < 0) { fclose(f); return -1; }
	char *b = malloc((size_t)n + 1);
	if (b == NULL) { fclose(f); return -1; }
	if (fread(b, 1, (size_t)n, f) != (size_t)n) { free(b); fclose(f); return -1; }
	b[n] = '\0';
	fclose(f);
	*outbuf = b;
	*outlen = (u64)n;
	return 0;
}

static int
export_owner_matches(const char *buf, u64 len, const char *path)
{
	static const char prefix[] = "//ww:module ";
	size_t pn = strlen(path);
	size_t need = sizeof prefix - 1 + pn + 1;
	return len >= need
	    && memcmp(buf, prefix, sizeof prefix - 1) == 0
	    && memcmp(buf + sizeof prefix - 1, path, pn) == 0
	    && buf[need - 1] == '\n';
}

struct importin {
	const char *path;
	const char *file;
	char *buf;
	u64 len;
	Node *ast;
};

struct importmap {
	const char *source;
	const char *path;
	int seen;
};

static Node *
parseinput(Arena *a, const char *file, char *buf, u64 len,
    const char *mod, const char *testsupport, int commandpackage, int *bad)
{
	Lex l;
	Parser p;
	lexinit(&l, a, file, buf, len);
	parserinit(&p, a, &l);
	p.testmodule = testsupport;
	p.commandpackage = commandpackage;
	if (mod != NULL) {
		p.pathmod = mod;
		p.curmod = mod;
	}
	Node *f = parsefile(&p);
	*bad = l.errs || p.errs;
	return f;
}

static int
canonical_pkgname(const char *path, const char *name)
{
	static const char hex[] = "0123456789abcdef";
	const char *prefix = "__wwi_";
	for (int i = 0; prefix[i]; i++)
		if (*name++ != prefix[i]) return 0;
	for (const unsigned char *p = (const unsigned char *)path; *p; p++) {
		if (*name++ != hex[*p >> 4] || *name++ != hex[*p & 15])
			return 0;
	}
	return *name == '\0';
}

static void
consider_pkgname(const char *path, const char *candidate,
	const char **name, const char **placeholder, int *conflict)
{
	if (canonical_pkgname(path, candidate)) {
		*placeholder = candidate;
		return;
	}
	if (*name != NULL && strcmp(*name, candidate) != 0) {
		*conflict = 1;
		return;
	}
	*name = candidate;
}

/* A paired interface owns PATH's declared name. Compiler-private names keep
 * transitive fact sections semantic and are ignored when a real name exists. */
static const char *
import_pkgname(struct importin *imports, int nimports, const char *path,
    Node *primary, int *conflict)
{
	const char *name = NULL;
	const char *placeholder = NULL;
	for (int i = 0; i < nimports; i++) {
		if (strcmp(imports[i].path, path) != 0) continue;
		Node *file = imports[i].ast;
		for (Node *p = file ? file->body : NULL; p; p = p->next) {
			if (p->module == NULL || p->pkgname == NULL
			    || strcmp(p->module, path) != 0)
				continue;
			consider_pkgname(path, p->pkgname, &name, &placeholder,
			    conflict);
			if (*conflict) return NULL;
		}
	}
	for (Node *p = primary ? primary->body : NULL; p; p = p->next) {
		if (p->module == NULL || p->pkgname == NULL
		    || strcmp(p->module, path) != 0)
			continue;
		consider_pkgname(path, p->pkgname, &name, &placeholder, conflict);
		if (*conflict) return NULL;
	}
	return name ? name : placeholder;
}

static int
bind_import_names(Node *list, struct importin *imports, int nimports,
    Node *primary, const char *testsupport)
{
	for (Node *u = list; u; u = u->next) {
		if (u->kind != N_USE || u->usepath == NULL) continue;
		/* The reserved test-support spelling is a compiler-owned alias, not
		 * source default-import syntax. */
		if (testsupport != NULL && strcmp(testsupport, "__wwtest") == 0
		    && strcmp(u->usepath, testsupport) == 0)
			continue;
		int conflict = 0;
		const char *name = import_pkgname(imports, nimports, u->usepath,
		    primary, &conflict);
		if (conflict) {
			fprintf(stderr,
			    "w6c: package %s has conflicting declared names in export data\n",
			    u->usepath);
			return -1;
		}
		if (name != NULL) {
			u->usepkgname = name;
			if (!u->useblank) {
				u->str = u->usealias ? u->usealias : name;
				u->strlen = strlen(u->str);
			}
		} else if (!u->imported) {
			fprintf(stderr,
			    "w6c: import %s has no declared package name in direct export data\n",
			    u->usepath);
			return -1;
		} else {
			/* A closure-only import need not contribute declarations to this
			 * interface.  Keep it canonical-path keyed without reinstalling
			 * the historical path-leaf qualifier. */
			if (!u->useblank) {
				u->str = u->usealias ? u->usealias : u->usepath;
				u->strlen = strlen(u->str);
			}
		}
	}
	return 0;
}

static void
appendnodes(Node **head, Node **tail, Node *list)
{
	if (list == NULL) return;
	if (*head == NULL) *head = list;
	else (*tail)->next = list;
	while (list->next != NULL) list = list->next;
	*tail = list;
}

int
main(int argc, char **argv)
{
	const char *src = NULL;
	const char *out = NULL;
	const char *wwiout = NULL;	/* -I <out.wwi>: M2 export-data producer */
	const char *testsupport = NULL;
	const char **testtargets = calloc((size_t)argc, sizeof *testtargets);
	const char *packageinit = NULL;
	const char *initdispatch = NULL;
	int testmode = 0;
	int testpackage = 0;
	int commandpackage = 0;
	int entrymode = 0;
	int sepmode = 0;		/* -c: #22 M3 separate-compile / primary-
				 * only codegen (emit imported==0 decls
				 * only; treat `.wwi` deps as external) */
	struct importin *imports = calloc((size_t)argc, sizeof *imports);
	struct importmap *maps = calloc((size_t)argc, sizeof *maps);
	if (imports == NULL || maps == NULL || testtargets == NULL) {
		fputs("w6c: out of memory\n", stderr);
		free(testtargets);
		free(maps);
		free(imports);
		return 1;
	}
	int nimports = 0;
	int nmaps = 0;
	int ntesttargets = 0;
	for (int i = 1; i < argc; i++) {
		const char *a = argv[i];
		if (strcmp(a, "-o") == 0 && i + 1 < argc) {
			out = argv[++i];
		} else if (strcmp(a, "-I") == 0 && i + 1 < argc) {
			wwiout = argv[++i];
		} else if (strcmp(a, "-T") == 0) {
			testmode = 1;
		} else if (strcmp(a, "--test-package") == 0) {
			testpackage = 1;
		} else if (strcmp(a, "--command-package") == 0) {
			commandpackage = 1;
		} else if (strcmp(a, "--entry") == 0) {
			entrymode = 1;
		} else if (strcmp(a, "--package-init-symbol") == 0) {
			if (i + 1 >= argc) {
				fputs("w6c: --package-init-symbol requires arg\n", stderr);
				return 2;
			}
			packageinit = argv[++i];
		} else if (strcmp(a, "--init-dispatch-symbol") == 0) {
			if (i + 1 >= argc) {
				fputs("w6c: --init-dispatch-symbol requires arg\n", stderr);
				return 2;
			}
			initdispatch = argv[++i];
		} else if (strcmp(a, "--test-support-module") == 0) {
			if (i + 1 >= argc) {
				fputs("w6c: --test-support-module requires arg\n", stderr);
				return 2;
			}
			testsupport = argv[++i];
		} else if (strcmp(a, "--test-target-package") == 0) {
			if (i + 1 >= argc) {
				fputs("w6c: --test-target-package requires arg\n", stderr);
				return 2;
			}
			testtargets[ntesttargets++] = argv[++i];
		} else if (strcmp(a, "-c") == 0) {
			sepmode = 1;
		} else if (strcmp(a, "--import") == 0) {
			if (i + 2 >= argc) {
				fputs("w6c: --import requires path and file\n", stderr);
				return 2;
			}
			imports[nimports].path = argv[++i];
			imports[nimports].file = argv[++i];
			nimports++;
		} else if (strcmp(a, "--import-map") == 0) {
			if (i + 2 >= argc) {
				fputs("w6c: --import-map requires source and path\n", stderr);
				return 2;
			}
			maps[nmaps].source = argv[++i];
			maps[nmaps].path = argv[++i];
			nmaps++;
		} else if (a[0] == '-') {
			fprintf(stderr, "w6c: unknown flag %s\n", a);
			return 2;
		} else if (src == NULL) {
			src = a;
		} else {
			fprintf(stderr, "w6c: only one input supported\n");
			return 2;
		}
	}
	if (src == NULL) {
		fputs("usage: w6c [-T|--test-package] [--command-package] [--entry] [--package-init-symbol symbol] [--init-dispatch-symbol symbol] [--test-target-package path] [-c] [-I out.wwi] "
		    "[--import path dep.wwi]... [--import-map source path]... [-o out.s] file.ww\n", stderr);
		return 2;
	}
	if (out != NULL && wwiout != NULL && strcmp(out, wwiout) == 0) {
		fputs("w6c: assembly and interface outputs must be distinct\n", stderr);
		return 2;
	}
	if (nimports > 0 && !sepmode) {
		fputs("w6c: --import requires -c\n", stderr);
		return 2;
	}
	if (nmaps > 0 && !sepmode) {
		fputs("w6c: --import-map requires -c\n", stderr);
		return 2;
	}
	if ((entrymode || testpackage || commandpackage) && !sepmode) {
		fputs("w6c: --entry, --test-package, and --command-package require -c\n", stderr);
		return 2;
	}
	if ((packageinit != NULL || initdispatch != NULL) && !sepmode) {
		fputs("w6c: package initialization symbols require -c\n", stderr);
		return 2;
	}
	if (packageinit != NULL && packageinit[0] == '\0') {
		fputs("w6c: --package-init-symbol is empty\n", stderr);
		return 2;
	}
	if (initdispatch != NULL && (initdispatch[0] == '\0' || !entrymode)) {
		fputs("w6c: --init-dispatch-symbol requires --entry and a non-empty symbol\n",
		    stderr);
		return 2;
	}
	if (testmode && testpackage) {
		fputs("w6c: -T and --test-package are mutually exclusive\n", stderr);
		return 2;
	}
	for (int i = 0; i < nimports; i++) {
		if (imports[i].path[0] == '\0') {
			fputs("w6c: --import path is empty\n", stderr);
			return 2;
		}
		if (i > 0 && strcmp(imports[i-1].path, imports[i].path) >= 0) {
			fputs("w6c: --import paths must be sorted and unique\n", stderr);
			return 2;
		}
	}
	for (int i = 0; i < nmaps; i++) {
		if (maps[i].source[0] == '\0' || maps[i].path[0] == '\0') {
			fputs("w6c: --import-map path is empty\n", stderr);
			return 2;
		}
		if (i > 0 && strcmp(maps[i-1].source, maps[i].source) >= 0) {
			fputs("w6c: --import-map sources must be sorted and unique\n",
			    stderr);
			return 2;
		}
		int direct = 0;
		for (int j = 0; j < nimports; j++)
			if (strcmp(maps[i].path, imports[j].path) == 0) {
				direct = 1;
				break;
			}
		if (!direct) {
			fputs("w6c: --import-map target is not a direct import\n",
			    stderr);
			return 2;
		}
	}
	if (testsupport != NULL && (!sepmode
	    || (strcmp(testsupport, "test") != 0
		&& strcmp(testsupport, "__wwtest") != 0))) {
		fputs("w6c: invalid --test-support-module\n", stderr);
		return 2;
	}
	for (int ti = 0; ti < ntesttargets; ti++) {
		const char *testtarget = testtargets[ti];
		if (!sepmode || !testmode || testtarget[0] == '\0') {
			fputs("w6c: invalid --test-target-package\n", stderr);
			return 2;
		}
		if (ti > 0 && strcmp(testtargets[ti - 1], testtarget) >= 0) {
			fputs("w6c: --test-target-package paths must be sorted and unique\n",
			    stderr);
			return 2;
		}
		int direct = 0;
		for (int i = 0; i < nimports; i++)
			if (strcmp(imports[i].path, testtarget) == 0)
				direct = 1;
		if (!direct) {
			fputs("w6c: --test-target-package is not a direct import\n",
			    stderr);
			return 2;
		}
	}

	Arena *a = newarena();
	Checker c;
	Cg cg;

	Node *head = NULL, *tail = NULL;
	for (int i = 0; i < nimports; i++) {
		if (slurp(imports[i].file, &imports[i].buf,
		    &imports[i].len) < 0) {
			fprintf(stderr, "w6c: import %s: cannot read %s\n",
			    imports[i].path, imports[i].file);
			return 1;
		}
		if (!export_owner_matches(imports[i].buf, imports[i].len,
		    imports[i].path)) {
			fprintf(stderr,
			    "w6c: import %s: export owner mismatch in %s\n",
			    imports[i].path, imports[i].file);
			return 1;
		}
		int bad = 0;
		Node *f = parseinput(a, imports[i].file, imports[i].buf,
		    imports[i].len, imports[i].path, testsupport, 0, &bad);
		if (bad) return 1;
		imports[i].ast = f;
		if (bind_import_names(f->list, imports, i + 1, NULL,
		    testsupport) < 0)
			return 1;
		appendnodes(&head, &tail, f->list);
	}

	char *buf;
	u64 len;
	if (slurp(src, &buf, &len) < 0) {
		fprintf(stderr, "w6c: %s: cannot read\n", src);
		return 1;
	}
	int bad = 0;
	/* --entry classifies an ordinary selected command.  The separate
	 * --command-package marker carries the same parser-only fact for command
	 * test variants whose code generation must not expose bare main. */
	Node *file = parseinput(a, src, buf, len, NULL, testsupport,
	    commandpackage || entrymode, &bad);
	if (bad) return 1;
	/* Vendor expansion changes only canonical identity. Source spelling and
	 * an optional file-local alias remain independent facts. */
	for (Node *u = file->list; u; u = u->next) {
		if (u->kind != N_USE || u->usepath == NULL) continue;
		const char *source = u->usesource ? u->usesource : u->usepath;
		for (int i = 0; i < nmaps; i++)
			if (strcmp(source, maps[i].source) == 0) {
				u->usepath = maps[i].path;
				maps[i].seen = 1;
				break;
			}
	}
	for (int i = 0; i < nmaps; i++) {
		if (maps[i].seen) continue;
		fputs("w6c: --import-map source is not in primary input\n", stderr);
		return 2;
	}
	/* Later direct interfaces may supply names for origin sections referenced
	 * by an earlier interface, so perform one complete metadata pass now. */
	for (int i = 0; i < nimports; i++)
		if (bind_import_names(imports[i].ast->list, imports, nimports,
		    NULL, testsupport) < 0)
			return 1;
	if (bind_import_names(file->list, imports, nimports, file,
	    testsupport) < 0)
		return 1;
	for (int ti = 0; ti < ntesttargets; ti++) {
		int seen = 0;
		for (Node *u = file->list; u; u = u->next) {
			if (u->kind != N_USE || u->imported || u->usepath == NULL
			    || strcmp(u->usepath, testtargets[ti]) != 0)
				continue;
			u->str = u->usepath;
			u->strlen = strlen(u->usepath);
			seen++;
		}
		if (seen != 1) {
			fputs("w6c: generated test target import is not unique\n",
			    stderr);
			return 2;
		}
	}
	if (head != NULL) {
		tail->next = file->list;
		file->list = head;
	}

	check_init(&c, a);
	c.is_test = testmode;
	c.is_test_package = testpackage;
	if (testsupport != NULL) c.test_module = testsupport;
	c.test_targets = testtargets;
	c.n_test_targets = ntesttargets;
	c.sep_mode = sepmode;
	c.package_init_symbol = packageinit;
	check_file(&c, file);
	if (c.errs) return 1;

	/* Anonymous streams keep every named destination untouched until checking,
	 * export writing, and lowering have all succeeded. A cgen fatal exits with
	 * only kernel-owned anonymous files open. */
	FILE *wf = NULL;
	if (wwiout) {
		wf = tmpfile();
		if (wf == NULL) {
			fprintf(stderr, "w6c: cannot stage %s\n", wwiout);
			return 1;
		}
		if (wwi_emit(&c, wf, file) != 0 || fflush(wf) != 0
		    || ferror(wf)) {
			fclose(wf);
			return 1;
		}
	}

	FILE *of = stdout;
	if (out) {
		of = tmpfile();
		if (of == NULL) {
			fprintf(stderr, "w6c: cannot stage %s\n", out);
			if (wf != NULL) fclose(wf);
			return 1;
		}
	}

	cg_init(&cg, a);
	cg.sep_mode = sepmode;
	/* Export production and entry identity are independent package-action
	 * properties. Legacy raw invocations without -I remain entry-like; every
	 * driver package now supplies -I, and only link roots add --entry. */
	cg.sep_isdep = (wwiout != NULL && !entrymode);
	cg.init_dispatch_symbol = initdispatch;
	cg_file(&cg, of, file);

	if (of != stdout && (fflush(of) != 0 || ferror(of))) {
		fclose(of);
		if (wf != NULL) fclose(wf);
		return 1;
	}
	struct publication pub[2] = {0};
	int npub = 0;
	if (wwiout != NULL) {
		pub[npub].dst = wwiout;
		if (materialize_stream(wf, &pub[npub]) < 0) goto publish_fail;
		npub++;
	}
	if (out != NULL) {
		pub[npub].dst = out;
		if (materialize_stream(of, &pub[npub]) < 0) goto publish_fail;
		npub++;
	}
	if (wf != NULL) fclose(wf);
	if (of != stdout) fclose(of);
	if (publish_all(pub, npub) < 0) goto publish_free_fail;
	for (int i = 0; i < npub; i++) publication_free(&pub[i]);
	freearena(a);
	for (int i = 0; i < nimports; i++) free(imports[i].buf);
	free(imports);
	free(maps);
	free(buf);
	return 0;

publish_fail:
	if (wf != NULL) fclose(wf);
	if (of != stdout) fclose(of);
	for (int i = 0; i < 2; i++) {
		if (pub[i].stage != NULL) (void)unlink(pub[i].stage);
		publication_free(&pub[i]);
	}
	return 1;
publish_free_fail:
	for (int i = 0; i < npub; i++) publication_free(&pub[i]);
	return 1;
}