#include "stdafx.h"
#include "Params.h"
#include "Debug.h"
#include "Compiler/Server/Main.h"
#include "Compiler/Exception.h"
#include "Compiler/Engine.h"
#include "Compiler/Package.h"
#include "Compiler/Repl.h"
#include "Compiler/Version.h"
#include "Core/Timing.h"
#include "Core/Io/StdStream.h"
#include "Core/Io/Text.h"
#include "Core/Convert.h"

// Wait a bit before exiting so that the user can read error messages?
// This is only applicable on Windows systems in cases where Storm was launched
// from a non-console environment.
bool waitBeforeExit = false;

void runRepl(Engine &e, const wchar_t *lang, Repl *repl) {
	TextInput *input = e.stdIn();
	TextOutput *output = e.stdOut();

	Str *line = null;
	while (true) {
		{
			StrBuf *prompt = new (e) StrBuf();
			if (line)
				*prompt << S("? ");
			else
				*prompt << lang << S("> ");

			output->write(prompt->toS());
			output->flush();
		}

		Str *data = input->readLine();
		if (!line) {
			line = data;
		} else {
			StrBuf *to = new (repl) StrBuf();
			*to << line << S("\n") << data;
			line = to->toS();
		}

		try {
			Repl::Result result = repl->eval(line, null);
			StrBuf *msg = new (repl) StrBuf();
			if (result.isSuccess()) {
				if (Str *r = result.result()) {
					*msg << S("=> ");
					msg->indentBy(new (e) Str(S("   ")));
					msg->indent();
					*msg << r;
					msg->dedent();
				}
			} else if (Str *e = result.isError()) {
				*msg << S("Error: ") << e;
			} else if (result.isTerminate()) {
				break;
			}

			if (!result.isIncomplete())
				line = null;

			output->writeLine(msg->toS());
		} catch (const storm::Exception *err) {
			output->writeLine(err->toS());
			line = null;
		} catch (const ::Exception &err) {
			std::wostringstream t;
			t << err << endl;
			output->writeLine(new (e) Str(t.str().c_str()));
			line = null;
		}
	}
}

int runRepl(Engine &e, Duration bootTime, Url *root, const wchar_t *lang, const wchar_t *input) {
	if (!input) {
		wcout << L"Welcome to the Storm compiler!" << endl;
		wcout << L"Root directory: " << root << endl;
		wcout << L"Compiler boot in " << bootTime << endl;
	}

	if (!lang)
		lang = L"bs";

	Name *replName = new (e) Name();
	replName->add(new (e) Str(S("lang")));
	replName->add(new (e) Str(lang));
	replName->add(new (e) Str(S("repl")));
	Function *replFn = as<Function>(e.scope().find(replName));
	if (!replFn) {
		wcerr << L"Could not find a repl for " << lang << L": no function named " << replName << L" exists." << endl;
		waitBeforeExit = true;
		return 1;
	}

	Value replVal(Repl::stormType(e));
	if (!replVal.mayReferTo(replFn->result)) {
		wcerr << L"The function " << replName << L" must return a subtype of "
			  << replVal << L", not " << replFn->result << endl;
		waitBeforeExit = true;
		return 1;
	}

	typedef Repl *(CODECALL *CreateRepl)();
	CreateRepl createRepl = (CreateRepl)replFn->ref().address();
	Repl *repl = (*createRepl)();

	if (input) {
		Repl::Result result = repl->eval(new (e) Str(input), null);
		if (result.isSuccess()) {
			if (Str *r = result.result()) {
				wcout << r << endl;
			}
		} else if (Str *e = result.isError()) {
			wcerr << L"Error: " << e << endl;
			waitBeforeExit = true;
			return 1;
		} else if (result.isIncomplete()) {
			wcerr << L"The input given to the REPL is incomplete." << endl;
			waitBeforeExit = true;
			return 1;
		}
	} else {
		repl->greet();
		runRepl(e, lang, repl);
	}

	return 0;
}

bool hasColon(const wchar_t *function) {
	for (; *function; function++)
		if (*function == ':')
			return true;
	return false;
}

int runFunction(Engine &e, Function *fn) {
	Value r = fn->result;
	if (!r.returnInReg()) {
		wcout << L"Found a main function: " << fn->identifier() << L", but it returns a value-type." << endl;
		wcout << L"This is not currently supported. Try running it through the REPL instead." << endl;
		waitBeforeExit = true;
		return 1;
	}

	// We can ignore the return value. But if it is an integer, we capture it and return it.
	code::Primitive resultType;
	if (code::PrimitiveDesc *resultDesc = as<code::PrimitiveDesc>(r.desc(e)))
		resultType = resultDesc->v;

	Nat resultSize = 0;
	if (resultType.kind() == code::primitive::integer)
		resultSize = resultType.size().current();

	RunOn run = fn->runOn();
	const void *addr = fn->ref().address();
	if (run.state == RunOn::named) {
		if (resultSize == 1) {
			os::FnCall<Byte, 1> call = os::fnCall();
			os::Future<Byte> result;
			os::Thread on = run.thread->thread()->thread();
			os::UThread::spawn(addr, false, call, result, &on);
			return result.result();
		} else if (resultSize == 4) {
			os::FnCall<Int, 1> call = os::fnCall();
			os::Future<Int> result;
			os::Thread on = run.thread->thread()->thread();
			os::UThread::spawn(addr, false, call, result, &on);
			return result.result();
		} else if (resultSize == 8) {
			os::FnCall<Long, 1> call = os::fnCall();
			os::Future<Long> result;
			os::Thread on = run.thread->thread()->thread();
			os::UThread::spawn(addr, false, call, result, &on);
			return (int)result.result();
		} else {
			os::FnCall<void, 1> call = os::fnCall();
			os::Future<void> result;
			os::Thread on = run.thread->thread()->thread();
			os::UThread::spawn(addr, false, call, result, &on);
			result.result();
		}
	} else {
		if (resultSize == 1) {
			typedef Byte (*Fn1)();
			Fn1 p = (Fn1)addr;
			return (*p)();
		} else if (resultSize == 4) {
			typedef Int (*Fn4)();
			Fn4 p = (Fn4)addr;
			return (*p)();
		} else if (resultSize == 8) {
			typedef Long (*Fn8)();
			Fn8 p = (Fn8)addr;
			return (int)(*p)();
		} else {
			typedef void (*Fn)();
			Fn p = (Fn)addr;
			(*p)();
		}
	}
	return 0;
}

int runFunction(Engine &e, const wchar_t *function) {
	SimpleName *name = parseSimpleName(new (e) Str(function));
	Named *found = e.scope().find(name);
	if (!found) {
		wcerr << L"Could not find " << function << endl;
		if (hasColon(function))
			wcerr << L"Did you mean to use '.' instead of ':'?" << endl;
		waitBeforeExit = true;
		return 1;
	}

	Function *fn = as<Function>(found);
	if (!fn) {
		wcout << function << L" is not a function." << endl;
		waitBeforeExit = true;
		return 1;
	}

	return runFunction(e, fn);
}

bool tryRun(Engine &e, Array<Package *> *pkgs, int &result) {
	SimplePart *mainPart = new (e) SimplePart(new (e) Str(S("main")));

	for (Nat i = 0; i < pkgs->count(); i++) {
		Package *pkg = pkgs->at(i);
		Function *main = as<Function>(pkg->find(mainPart, Scope(pkg)));
		if (!main)
			continue;

		result = runFunction(e, main);
		return true;
	}

	// Failed.
	return false;
}

int runTests(Engine &e, const wchar_t *package, bool recursive) {
	waitBeforeExit = true;

	SimpleName *name = parseSimpleName(new (e) Str(package));
	Named *found = e.scope().find(name);
	if (!found) {
		wcerr << L"Could not find " << package << endl;
		return 1;
	}

	Package *pkg = as<Package>(found);
	if (!pkg) {
		wcerr << package << L" is not a package." << endl;
		return 1;
	}

	SimpleName *testMain = parseSimpleName(e, S("test.runCmdline"));
	testMain->last()->params->push(Value(StormInfo<Package>::type(e)));
	testMain->last()->params->push(Value(StormInfo<Bool>::type(e)));
	Function *f = as<Function>(e.scope().find(testMain));
	if (!f) {
		wcerr << L"Could not find the test implementation's main function: " << testMain << endl;
		return 1;
	}

	if (f->result.type != StormInfo<Bool>::type(e)) {
		wcerr << L"The signature of the test implementation's main function ("
			  << testMain << L") is incorrect. Expected Bool, but got "
			  << f->result.type->identifier() << endl;
		return 1;
	}

	os::FnCall<Bool> c = os::fnCall().add(pkg).add(recursive);
	Bool ok = c.call(f->ref().address(), false);
	return ok ? 0 : 1;
}

// Returns array of paths that we should try to run if they exist.
Array<Package *> *importPkgs(Engine &into, const Params &p) {
	Array<Package *> *result = new (into) Array<Package *>();

	for (Nat i = 0; i < p.import.size(); i++) {
		const Import &import = p.import[i];

		Url *path = parsePath(new (into) Str(import.path));
		if (!path->absolute())
			path = cwdUrl(into)->push(path);

		// Update file/directory flag.
		path = path->updated();

		if (!path->exists()) {
			wcerr << L"WARNING: The path " << *path << L" could not be found. Will not import it." << endl;
			continue;
		}

		SimpleName *n = null;
		if (import.into) {
			n = parseSimpleName(new (into) Str(import.into));
		} else {
			n = parseSimpleName(path->title());
		}

		if (n->empty())
			continue;

		NameSet *ns = into.nameSet(n->parent(), true);
		Package *pkg = new (into) Package(n->last()->name, path);

		try {
			ns->add(pkg);
		} catch (storm::TypedefError *) {
			wcerr << L"WARNING: Unable to import the package " << path << L" as " << n
				  << L" since a package with the same name already exists.";

			if (import.into) {
				wcerr << L"Try specifying a different name." << endl;
			} else {
				// TODO? If import.tryRun, maybe try an anonymous name?
				wcerr << L"Either rename the file or directory to a different name, or use the -I flag to specify a different name." << endl;
			}

			continue;
		}

		if (import.tryRun)
			result->push(pkg);
	}

	return result;
}

void showVersion(Engine &e) {
	waitBeforeExit = true;

	Array<VersionTag *> *tags = storm::versions(e.package(S("core")));
	if (tags->empty()) {
		wcerr << L"Error: No version information found. Make sure this release was compiled correctly." << std::endl;
		return;
	}

	VersionTag *best = tags->at(0);
	for (Nat i = 1; i < tags->count(); i++) {
		if (best->path()->count() > tags->at(i)->path()->count())
			best = tags->at(i);
	}

	wcout << best->version << std::endl;
}

void createArgv(Engine &e, const vector<const wchar_t *> &argv) {
	Array<Str *> *out = new (e) Array<Str *>();
	out->reserve(Nat(argv.size()));
	for (size_t i = 0; i < argv.size(); i++) {
		*out << new (e) Str(argv[i]);
	}
	e.argv(out);
}

struct RootUrlData {
	Params &params;
	Url *created;
};

static Url *createRootUrl(Engine &e, void *d) {
	RootUrlData *data = (RootUrlData *)d;

	if (const wchar_t *fromParams = data->params.root) {
#if defined(WINDOWS)
		data->created = parsePathAsDir(e, fromParams);
#else
		data->created = parsePathAsDir(new (e) Str(toWChar(e, fromParams)));
#endif
		if (data->created->relative())
			data->created = data->created->makeAbsolute(cwdUrl(e));
	} else {
#if defined(DEBUG)
		data->created = dbgRootUrl(e)->push(new (e) Str(S("root")));
#elif defined(STORM_ROOT)
		data->created = parsePathAsDir(e, S(SHORT_STRING(STORM_ROOT)));
#else
		data->created = executableUrl(e)->push(new (e) Str(S("root")));
#endif
	}

	return data->created;
}

int stormMain(int argc, const wchar_t *argv[], void *stackBase) {
	Params p(argc, argv);

	switch (p.mode) {
	case Params::modeHelp:
		help(argv[0]);
		waitBeforeExit = true;
		return 0;
	case Params::modeError:
		wcerr << L"Error in parameters: " << p.modeParam;
		if (p.modeParam2)
			wcerr << p.modeParam2;
		wcerr << endl;
		waitBeforeExit = true;
		return 1;
	}

#ifdef WINDOWS
	if (p.noConsole) {
		FreeConsole();
	}
#endif

	Moment start;

	// Start an Engine.
	RootUrlData root = { p, null };
	Engine e(&createRootUrl, &root, Engine::reuseMain, stackBase);
	Moment end;

	try {

		if (!p.argv.empty())
			createArgv(e, p.argv);

		Array<Package *> *runPkgs = importPkgs(e, p);

		int result = 1;

		try {
			switch (p.mode) {
			case Params::modeAuto:
				if (!tryRun(e, runPkgs, result)) {
					// If none of them contained a main function, launch the repl.
					result = runRepl(e, (end - start), root.created, null, null);
				}
				break;
			case Params::modeRepl:
				result = runRepl(e, (end - start), root.created, p.modeParam, p.modeParam2);
				break;
			case Params::modeFunction:
				result = runFunction(e, p.modeParam);
				break;
			case Params::modeTests:
				result = runTests(e, p.modeParam, false);
				break;
			case Params::modeTestsRec:
				result = runTests(e, p.modeParam, true);
				break;
			case Params::modeVersion:
				showVersion(e);
				result = 0;
				break;
			case Params::modeServer:
				server::run(e, proc::in(e), proc::out(e));
				result = 0;
				break;
			default:
				throw new (e) InternalError(S("Unknown mode."));
			}
		} catch (const storm::Exception *e) {
			wcerr << e << endl;
			waitBeforeExit = true;
			result = 1;
			// Fall-thru to wait for UThreads.
		} catch (const ::Exception &e) {
			// Sometimes, we need to print the exception before the engine is destroyed.
			wcerr << e << endl;
			waitBeforeExit = true;
			return 1;
		}

		// Allow 1 s for all UThreads on the Compiler thread to terminate.
		Moment waitStart;
		while (os::UThread::leave() && Moment() - waitStart > time::s(1))
			;

		return result;

	} catch (storm::Exception *e) {
		wcerr << L"Unhandled exception:\n" << *e << endl;
		return 1;
	}
}

#ifdef WINDOWS

const UINT targetCP = 65001;

int _tmain(int argc, const wchar *argv[]) {
	int result = 0;

	// Set output codepage to UTF-8
	UINT inputCP = GetConsoleCP();
	UINT outputCP = GetConsoleOutputCP();
	if (inputCP != 0 && inputCP != targetCP)
		SetConsoleCP(targetCP);
	if (outputCP != 0 && outputCP != targetCP)
		SetConsoleOutputCP(targetCP);

	try {
		result = stormMain(argc, argv, &argv);
	} catch (const ::Exception &e) {
		wcerr << L"Unhandled exception: " << e << endl;
		result = 1;
	}

	if (waitBeforeExit) {
		// If the console would disappear when we exit, wait a bit to let the user read any messages
		// in it.
		DWORD entries[2];
		DWORD count = GetConsoleProcessList(entries, 2);

		if (count == 1) {
			// Only we (note, 'count' may be zer if we detached from it):
			wcerr << L"\nPress any key to exit" << endl;
			_getwch();
		}
	}

	// Restore codepage:
	if (inputCP != 0 && inputCP != targetCP)
		SetConsoleCP(inputCP);
	if (outputCP != 0 && outputCP != targetCP)
		SetConsoleOutputCP(outputCP);

	return result;
}

#else

int main(int argc, const char *argv[]) {
	try {
		vector<String> args(argv, argv+argc);
		vector<const wchar_t *> c_args(argc);
		for (int i = 0; i < argc; i++)
			c_args[i] = args[i].c_str();

		return stormMain(argc, &c_args[0], &argv);
	} catch (const ::Exception &e) {
		wcerr << "Unhandled exception: " << e << endl;
		return 1;
	}
}

#endif