#include "stdafx.h"
#include "Package.h"
#include "Core/StrBuf.h"
#include "Core/Str.h"
#include "Core/Set.h"
#include "Core/Io/Text.h"
#include "Core/Set.h"
#include "Gc/ObjMap.h"
#include "Engine.h"
#include "Reader.h"
#include "Exception.h"
#include "TemplateUpdate.h"
#include "ActiveFunctions.h"

namespace storm {

	Package::Package(Str *name) : NameSet(name), pkgPath(null), loading(null), exportedLoaded(false) {
		// The default is to discard source automatically.
		NameSet::discardSource();
	}

	Package::Package(Url *path) : NameSet(path->name()), pkgPath(path), loading(null), exportedLoaded(false) {

		// The default is to discard source automatically.
		NameSet::discardSource();

		engine().pkgMap()->put(pkgPath, this);

		documentation = new (this) PackageDoc(this);
	}

	Package::Package(Str *name, Url *path) : NameSet(name), pkgPath(path), loading(null), exportedLoaded(false) {
		// The default is to discard source automatically.
		NameSet::discardSource();

		engine().pkgMap()->put(pkgPath, this);

		documentation = new (this) PackageDoc(this);
	}

	NameLookup *Package::parent() const {
		// We need to be able to return null.
		return parentLookup();
	}

	Url *Package::url() const {
		return pkgPath;
	}

	void Package::setUrl(Url *url) {
		STORM_DBG_ASSERT_S(!engine().has(bootDone), "Shall not be done after boot is complete!");
		if (pkgPath)
			engine().pkgMap()->remove(pkgPath);
		pkgPath = url;
		engine().pkgMap()->put(pkgPath, this);

		if (pkgPath->dir()) {
			for (Iter i = begin(), e = end(); i != e; ++i) {
				if (Package *p = as<Package>(i.v())) {
					if (p->params->count() > 0)
						continue;
					Url *sub = url->pushDir(p->name);
					if (!sub->exists())
						continue;

					p->setUrl(sub);
				}
			}
		}

		if (!documentation)
			documentation = new (this) PackageDoc(this);
	}

	MAYBE(Named *) Package::has(Named *item) const {
		if (loading)
			if (Named *found = loading->has(item))
				return found;

		return NameSet::has(item);
	}

	void Package::add(Named *item) {
		if (loading) {
			// If we should not allow duplicates, check for duplicates with ourself first. If this
			// would be a duplicate, try to add it to ourself to generate a suitable error message
			// (we know that 'add' will throw in this case).
			if (!loadingAllowDuplicates && NameSet::has(item))
				NameSet::add(item);

			loading->add(item);
			this->makeChild(item);
		} else {
			NameSet::add(item);
		}
	}

	void Package::add(Template *item) {
		if (loading) {
			// We cannot really check for duplicates for templates.
			loading->add(item);
		} else {
			NameSet::add(item);
		}
	}

	Bool Package::remove(Named *item) {
		if (loading) {
			if (loading->remove(item))
				return true;
		}

		return NameSet::remove(item);
	}

	Bool Package::remove(Template *item) {
		if (loading) {
			if (loading->remove(item))
				return true;
		}

		return NameSet::remove(item);
	}

	NameSet::Iter Package::begin() const {
		if (loading)
			return NameSet::begin(loading);
		else
			return NameSet::begin();
	}

	NameSet::Iter Package::end() const {
		if (loading)
			return loading->end();
		else
			return NameSet::end();
	}

	MAYBE(Named *) Package::find(SimplePart *part, Scope source) {
		if (loading)
			if (Named *found = loading->find(part, source))
				return found;

		if (Named *found = NameSet::find(part, source))
			return found;

		return null;
	}

	Bool Package::loadName(SimplePart *part) {
		// We're only loading packages this way.
		if (part->params->empty()) {
			if (Package *pkg = loadPackage(part->name)) {
				add(pkg);
				return true;
			}
		}

		// There may be other things to load!
		return false;
	}

	Bool Package::loadAll() {
		// Do not load things before the compiler has properly started. As everything is a chaos at
		// that time, loading anything external will most likely fail. So we retry later.
		if (!engine().has(bootDone))
			return false;

		// Nothing to load.
		if (!pkgPath)
			return true;

		// Load exports if needed.
		loadExports();

		Array<Url *> *files = new (this) Array<Url *>();

		if (pkgPath->dir()) {
			// A directory, the normal case:
			Array<Url *> *children = pkgPath->children();

			// Load any remaining packages:
			for (Nat i = 0; i < children->count(); i++) {
				Url *now = children->at(i);
				if (now->dir()) {
					Str *name = now->name();
					if (!tryFind(new (this) SimplePart(name), Scope()))
						add(loadPackage(name));
				} else {
					files->push(now);
				}
			}
		} else {
			// It is a file. Just load it. This is used when specifying files on the command line for example.
			files->push(pkgPath);
		}

		// Load all code.
		loadFiles(files);

		return true;
	}

	void Package::addExport(Package *pkg) {
		if (!exported)
			exported = new (this) Array<Package *>();

		// Remove duplicates: generally quite cheap to do, and prevents lookups from being expensive.
		for (Nat i = 0; i < exported->count(); i++)
			if (exported->at(i) == pkg)
				return;

		exported->push(pkg);
	}

	Array<Package *> *Package::exports() {
		loadExports();

		if (exported)
			return new (this) Array<Package *>(*exported);
		else
			return new (this) Array<Package *>();
	}

	Array<Package *> *Package::recursiveExports() {
		Array<Package *> *result = exports();
		Set<Package *> *seen = new (this) Set<Package *>();

		// Note: 'result->count()' will change.
		for (Nat at = 0; at < result->count(); at++) {
			Package *examine = result->at(at);

			examine->loadExports();
			if (examine->exported) {
				for (Nat i = 0; i < examine->exported->count(); i++) {
					Package *add = examine->exported->at(i);
					if (!seen->has(add)) {
						seen->put(add);
						result->push(add);
					}
				}
			}
		}

		return result;
	}

	void Package::retainSource() {
		stopDiscardSource();
	}

	void Package::discardSource() {
		// We don't need to propagate this message, we emit it ourselves.
	}

	void Package::toS(StrBuf *to) const {
		if (parent())
			*to << L"Package " << identifier();
		else
			*to << L"Root package (" << name << L")";

		if (pkgPath)
			*to << L"(in " << pkgPath << L")";
		else
			*to << L"(virtual)";

		if (false) {
			*to << L"\n";
			Indent z(to);
			NameSet::toS(to);
		}
	}

	void Package::loadExports() {
		if (exportedLoaded)
			return;
		exportedLoaded = true;

		if (!pkgPath)
			return;

		Url *exports = pkgPath->push(new (this) Str(S("export")));
		if (!exports->exists())
			return;

		if (!exported)
			exported = new (this) Array<Package *>();

		Scope root = engine().scope();
		TextInput *input = readText(exports);
		SrcPos pos(exports, 0, 0);
		while (input->more()) {
			Str *line = input->readLine();

			// Update position. We only count newlines as 1 character, regardless of how they are
			// represented in the file.
			if (pos.end != 0)
				pos.start = pos.end + 1;
			pos.end = pos.start + line->peekLength(); // Not always exact, but good enough.

			// Trim whitespace to make "parsing" easier.
			line = trimWhitespace(line);

			// Empty line, ignore.
			if (line->begin() == line->end())
				continue;

			// Comment, ignore.
			if (line->begin().v() == Char('#'))
				continue;
			if (line->begin().v() == Char('/') && (line->begin() + 1).v() == Char('/'))
				continue;

			Name *name = parseComplexName(line);
			if (!name)
				throw new (this) SyntaxError(pos, TO_S(this, "Invalid format of names. Use: 'x.y.z(a, b.c)'. Comments start with '#' or '//'"));

			Named *found = root.find(name);
			if (!found)
				throw new (this) SyntaxError(pos, TO_S(this, S("The name ") << name << S(" does not refer to anything.")));

			Package *p = as<Package>(found);
			if (!p)
				throw new (this) SyntaxError(pos, TO_S(this, S("The name ") << found << S(" does not refer to a package.")));

			exported->push(p);
		}
	}

	Package *Package::loadPackage(Str *name) {
		// Virtual package, can not be auto loaded.
		if (!pkgPath)
			return null;

		// If pkgPath was a file, then we can't auto-load any sub-packages.
		if (!pkgPath->dir()) {
			return null;
		}

		// Name parts that are empty do not play well with the Url.
		if (name->empty())
			return null;

		// TODO: Make sure this is case sensitive!
		Url *sub = pkgPath->pushDir(name);
		if (!sub->exists())
			return null;

		return new (this) Package(sub);
	}

	void Package::loadFiles(Array<Url *> *files) {
		// Load everything into a separate NameSet so that we can easily roll back in case of problems.
		loading = new (this) NameSet(name, params);
		loadingAllowDuplicates = false;

		try {
			Map<Str *, PkgFiles *> *readers = readerName(files);
			Array<PkgReader *> *load = createReaders(readers);

			// Load everything!
			read(load);

			// All is well, merge with ourselves...
			// Note: Source is discarded when merge it here.
			NameSet::merge(loading);
			loading = null;

		} catch (...) {
			// Find and discard any generated types that were produced during the load and discard them.
			removeTemplatesFrom(loading);
			// Discard the partially loaded results.
			loading = null;
			throw;
		}
	}

	static Str *noReaderWarning(SimpleName *name, Array<Url *> *files) {
		StrBuf *msg = new (name) StrBuf();
		*msg << L"No reader for [";
		Url *rootUrl = name->engine().package()->url();
		for (nat i = 0; i < files->count(); i++) {
			if (i > 0)
				*msg << L", ";
			*msg << files->at(i)->relative(rootUrl);
		}
		*msg << L"] (should be " << name << L")";

		return msg->toS();
	}

	Array<PkgReader *> *Package::createReaders(Map<Str *, PkgFiles *> *readers) {
		typedef Map<Str *, PkgFiles *> ReaderMap;
		Array<PkgReader *> *r = new (this) Array<PkgReader *>();
		SimpleName *me = path();

		SimpleName *delayName = null;
		PkgFiles *delayFiles = null;

		for (ReaderMap::Iter i = readers->begin(), end = readers->end(); i != end; ++i) {
			SimpleName *name = i.v()->name;

			// Load ourselves last.
			if (name->parent()->sameAs(me)) {
				delayName = name;
				delayFiles = i.v();
				continue;
			}

			PkgReader *reader = createReader(name, i.v()->files, this);
			if (reader) {
				r->push(reader);
			} else {
				WARNING(noReaderWarning(name, i.v()->files)->c_str());
			}
		}

		if (delayName && delayFiles) {
			PkgReader *reader = createReader(delayName, delayFiles->files, this);
			if (reader)
				r->push(reader);
		}

		return r;
	}

	void Package::reload() {
		if (!pkgPath)
			return;

		Array<Url *> *files = new (this) Array<Url *>();
		Array<Url *> *all = pkgPath->children();
		for (Nat i = 0; i < all->count(); i++)
			if (!all->at(i)->dir())
				files->push(all->at(i));

		reload(files, true);
	}

	void Package::reload(Url *file) {
		reload(new (this) Array<Url *>(file));
	}

	void Package::reload(Array<Url *> *files) {
		reload(files, false);
	}

	// Context for diffing entities during a reload.
	class ReloadDiff : public NameDiff {
	public:
		ReloadDiff(ReplaceContext *ctx, Array<Url *> *files, Bool complete) : ctx(ctx) {
			removeItems = new (files) Array<Named *>();
			removeTemplates = new (files) Array<Template *>();
			update = new (files) Array<NamedPair>();
			this->files = new (files) Set<Url *>();

			for (Nat i = 0; i < files->count(); i++)
				this->files->put(files->at(i));
		}

		// We don't need to worry about new things.
		virtual void added(Named *item) {}
		virtual void added(Template *) {}

		// Old things, we just need to keep track of, so that we can remove them later.
		virtual void removed(Named *item) {
			if (handled(item->pos))
				removeItems->push(item);
		}
		virtual void removed(Template *item) {
			TODO(L"Handle templates(generators)! (they don't have a position)");
			removeTemplates->push(item);
		}

		virtual void changed(Named *old, Named *changed) {
			if (!handled(old->pos)) {
				// We have a duplicate definition!
				throw new (old) TypedefError(
					changed->pos,
					TO_S(old, changed << S(" is already defined at:\n@") << old->pos << S(": here")));
			}

			// Sanity-check the replacement now.
			if (Str *msg = changed->canReplace(old, ctx)) {
				// TODO: check if it makes sense to replace the entity completely. We should at
				// least try to invalidate some usages of the entity if we remove it!
				PLN(L"WARNING: " << changed->pos << L": " << msg << L" Replacing the entity.");
				// There was an issue with replacing the entity. See if we can remove the old one
				// and insert the new one instead!
				removeItems->push(old);
			} else {
				// Make the change happen.
				update->push(NamedPair(old, changed));
			}
		}

		// Context.
		ReplaceContext *ctx;

		// Things we need to remove from the package before merging.
		Array<Named *> *removeItems;
		Array<Template *> *removeTemplates;

		// Things to update. "old" -> "new".
		Array<NamedPair> *update;

		// Files we're examining. If null, we examine all files.
		Set<Url *> *files;

		// Is a particular entity handled by us? I.e., shall we treat it as if it exists?
		Bool handled(SrcPos pos) {
			if (!pos.file)
				return false;

			return !files || files->has(pos.file);
		}
	};

	void Package::reload(Array<Url *> *files, Bool complete) {
		// If we're not loaded yet, just load the package instead.
		if (!allLoaded()) {
			forceLoad();
			return;
		}

		NameSet *temporary = new (this) NameSet(name, params);
		loading = temporary;
		loadingAllowDuplicates = true;

		ReplaceContext *context = new (this) ReplaceContext();
		ReloadDiff diff(context, files, complete);

		try {
			// This part is very similar to 'loadFiles'.
			Map<Str *, PkgFiles *> *readers = readerName(files);
			Array<PkgReader *> *load = createReaders(readers);

			read(load);

			// Figure out which types in the old tree are equivalent to ones in the newly loaded
			// subtree. We need to disable lookups from 'temporary' while doing this as it would
			// otherwise find the new types instead of the old ones.
			loading = null;
			context->buildTypeEquivalence(this, temporary);
			loading = temporary;

			// Find templates that were generated while loading and replace them as well.
			Array<NamedPair> *templates = replaceTemplatesFrom(loading, context);
			// Put them inside 'diff' so that we perform the replacements down the line. We're
			// making sanity checks here as well.
			for (Nat i = 0; i < templates->count(); i++) {
				NamedPair p = templates->at(i);
				if (Str *msg = p.to->canReplace(p.from, context)) {
					StrBuf *m = new (this) StrBuf();
					*m << S("Unable to replace ") << p.from << S(": ") << msg;
					throw new (this) ReplaceError(p.to->pos, m->toS());
				}
			}
			// Do all replacements. Note: 'diff' is not yet touched, so we can replace its array.
			diff.update = templates;

			// Now, try to merge all entities inside 'loading', populating the 'diff' variable
			// with what to do after performing some sanity checks of the update operations.
			NameSet::diff(loading, diff, context);
		} catch (...) {
			// Remove any generated templates.
			removeTemplatesFrom(loading);
			// Discard the partially loaded results.
			loading = null;
			throw;
		}

		// Collect errors during the update process.
		Array<Exception *> *errors = new (this) Array<Exception *>();

		try {
			// Pause threads - we're going to do some weird things!
			// TODO: It might be a bit dangerous to pause threads this early, the reload process
			// might execute user code. It should be executing on the compiler thread, but might
			// involve other threads as well.
			PauseThreads pause(engine());
			ReplaceTasks *tasks = new (this) ReplaceTasks(pause, errors);

			// Once this process is started, we can't roll it back anymore, so keep going as far as
			// possible, regardless of whether certain steps fail or not.
			for (Nat i = 0; i < diff.update->count(); i++) {
				NamedPair p = diff.update->at(i);
				p.to->replace(p.from, tasks, context);

				// Remove the old one.
				NameSet *in = as<NameSet>(p.from->parent());
				if (in) {
					in->NameSet::remove(p.from);
				}
				// Note: If 'in' is not a NameSet, then 'p.from' is likely something anonymous, like
				// a lambda function, which is fine to not remove.
			}

			// Remove items to make the merge go smootly. TODO: We should check so that removed items
			// are not used at some point, perhaps earlier than this.
			for (Nat i = 0; i < diff.removeItems->count(); i++)
				NameSet::remove(diff.removeItems->at(i));
			for (Nat i = 0; i < diff.removeTemplates->count(); i++)
				NameSet::remove(diff.removeTemplates->at(i));

			// Rewrite memory according to the updates.
			// In particular, we need to:
			// - replace all occurrences of the old entities with the new ones (at least for types).
			// - modify the layout of any types requiring that.
			// - update active functions.
			tasks->apply();

			// Now, we're done. We have resolved all conflicts, so now it is safe to merge the two
			// NameSets. Note: this may discard source, which is needed to update active functions.
			NameSet::merge(loading);
			loading = null;
		} catch (...) {
			// This indicates an internal error while merging. We will most likely not be able to
			// recover from this gracefully.
			TODO(L"How can we recover here?");
			loading = null;
			throw;
		}

		// Throw errors if necessary:
		if (errors->any())
			throw new (this) MultiException(errors);
	}

	/**
	 * Documentation
	 */

	PackageDoc::PackageDoc(Package *owner) : pkg(owner) {}

	Doc *PackageDoc::get() {
		Doc *result = doc(pkg);

		if (Url *url = pkg->url()) {
			Url *file = url->push(new (this) Str(S("README")));
			if (file->exists())
				result->body = readAllText(file);
		}

		return result;
	}


	/**
	 * Helpers for Storm.
	 */

	MAYBE(Package *) package(Url *path) {
		return path->engine().package(path);
	}

	Package *rootPkg(EnginePtr e) {
		return e.v.package();
	}

}
