use ui;
use core:io;
use core:lang;
use core:sync;
use lang;
use lang:bs:macro;

/**
 * Class representing some loaded program that we want to execute and possibly single-step it.
 *
 * Keeps track of the current state of the threads, and allows accessing the global variables in the
 * program. The data visible here is essentially a direct representation of the currently running
 * program. It may thereby be updated at any time.
 */
class Program on Render {
	// Package created by us as a temporary workspace.
	private Package pkg;

	// All view hints that are active.
	public progvis:data:ViewHints[] hints;

	// Pointer to the 'main' function.
	public Fn<void> main;

	// Original state of global variables, so that we can restore them.
	private SavedGlobals? globals;

	// Source code, pre formatted.
	private Url->Text code;

	// Output from the program since the last check. This is not wired to stdout in any way at the
	// moment, but we could do that. Use 'addOutput' to access it, with convenience functions below.
	private StrBuf output;

	// Thread ID:s. We maintain our own thread ID:s that have lower numbers. They are only unique for each program.
	private Nat usedThreadId;

	// Last thread that terminated. We use this to tell the user which thread crashed. If multiple
	// threads crash at the same time, this might be inaccurate, but it is good enough for now.
	private Nat lastTerminatedId;

	// Are we currently terminating all threads? We don't send update notifications during this time.
	private Bool terminating;

	// All threads in the program. The key is the ID of the UThread that executes the code.
	public Word->ProgThread threads;

	// Sorted list of threads. Sorted by their thread ID.
	public ProgThread[] sortedThreads;

	// New allocations since start or last check.
	private unsafe:RawPtr[] leakableAllocs;

	// Event called whenever the state of a thread changed.
	private Fn<void, Bool>? changeNotification;

	// Event called whenever an error is encountered from one of the threads.
	private Fn<void, ThreadError>? errorNotification;

	// Event called when a particular thread has been advanced (for recording). The bool indicates
	// that a "large" step (i.e. to the next barrier) was taken.
	private Fn<void, Nat, Bool>? advanceNotification;

	// Set the change notification. The 'bool' parameter indicates wheter a true breakpoint was hit
	// or not, i.e. if the thread intends to sleep for a while after the notification. This is
	// equivalent to the question: Is it useful to redraw the screen in addition to checking for
	// errors?
	public assign onChange(Fn<void, Bool> fn) {
		changeNotification = fn;
		fn;
	}

	// Set the error notification.
	public assign onError(Fn<void, ThreadError> fn) {
		errorNotification = fn;
		fn;
	}

	// Set the "advance" notification. This is called whenever each thread is stepped.
	public assign onAdvance(Fn<void, Nat, Bool> fn) {
		advanceNotification = fn;
		fn;
	}

	// Call the notification. Parameter indicates whether a breakpoint was hit or not.
	void notifyChange(Bool breakpoint) {
		// Inhibit notifications if we're terminating.
		if (!terminating) {
			if (changeNotification)
				changeNotification.call(breakpoint);
		}
	}

	// Call the advance notification.
	void notifyAdvance(Nat thread, Bool barrier) {
		if (advanceNotification)
			advanceNotification.call(thread, barrier);
	}

	// Notify that there has been an error.
	void notifyProgramError(Nat thread, Exception exception) {
		if (errorNotification) {
			Str type = "exception";
			Str msg = exception.message;

			// See if we have more information.
			if (exception as progvis:TypedError)
				type = exception.type;

			errorNotification.call(ThreadError(type, msg, thread));
		} else {
			print("The program encountered an error:\n${exception}");
		}
	}

	package init(Package pkg, Fn<void> main, progvis:data:ViewHints[] hints) {
		init() {
			pkg = pkg;
			main = main;
			hints = hints;
		}

		// Add us into the global map.
		programInstances.put(this);

		// Save global variables.
		globals = SavedGlobals:save(findGlobals());
	}

	// Find supported file types.
	public Set<Str> supportedFiles() : static {
		findSupportedFiles();
	}

	public Program load(Url file) : static {
		load([file]);
	}

	public Program load(Url[] files) : static {
		loadProgram(files);
	}

	// Add output to the buffer. Note: We don't call 'onChange' for this, since the program will
	// soon be halted at the next step anyway, which causes a full repaint anyway.
	public void addOutput(Str data) {
		output << data;
	}

	// Get all output in the buffer. Returns null if no new output.
	public Str? getOutput() {
		if (output.empty)
			return null;

		Str result = output.toS;
		output.clear();
		result;
	}

	// Add a thread to our thread structure.
	private void addThread(Word id, ProgThread thread) {
		threads.put(id, thread);

		for (Nat i = sortedThreads.count; i > 0; i--) {
			if (sortedThreads[i-1].threadId < thread.threadId) {
				sortedThreads.insert(i, thread);
				return;
			}
		}

		sortedThreads.insert(0, thread);
	}

	// Remove a thread from our thread structure.
	private void removeThread(Word id, ProgThread thread) {
		// Manually clear the memory watchers now. Otherwise we will retain references too long in
		// some cases.
		thread.memory.clearAll();

		threads.remove(id);

		for (i, v in sortedThreads) {
			if (v is thread) {
				sortedThreads.remove(i);
				return;
			}
		}
	}
	private void removeThread(Word id) {
		var found = threads.find(id);
		if (found == threads.end)
			return;

		removeThread(id, found.v);
	}

	// Start a new thread executing a particular function assumed to be located inside the
	// environment specified earlier. Always spawns a new thread.
	void spawn(Fn<void> fn) {
		(spawn run(fn)).detach();
	}

	// Run the function.
	void run(Fn<void> fn) {
		try {
			// Add the thread now so that it may be used to capture the result of 'main'.
			addThread(currentUThread(), ProgThread(this, ++usedThreadId));

			fn.call();
		} catch (ExitError e) {
			// Just ignore it.
		} catch (Exception e) {
			notifyProgramError(lastTerminatedId, e);
		}

		// If the thread remains by now, remove it.
		// This is just to be safe. It should not happen.
		removeThread(currentUThread());
	}

	// Load source code.
	package void loadSource(Url[] files) {
		for (f in files) {
			code.put(f, highlightSource(f));
		}
	}

	// Get source for a file.
	public Text? source(Url file) {
		code.at(file);
	}

	// Find and return all global variables in the program managed by this instance.
	GlobalVar[] findGlobals() {
		findGlobalsHelp(pkg);
	}

	// Check if all threads are paused currently.
	Bool threadsPaused() {
		for (t in threads)
			if (!t.paused)
				return false;
		true;
	}

	// Check if all threads are waiting for user interaction.
	Bool threadsWaiting() {
		for (t in threads)
			if (!t.waiting)
				return false;
		true;
	}

	// Terminate any threads running in the program, and wait for them to exit properly.
	void terminate() {
		// Inhibit change notifications - we want to make it look like all threads terminated at the same time.
		terminating = true;

		for (t in threads) {
			t.terminate();
		}

		// Wait for them to terminate...
		while (threads.count > 0) {
			yield();
		}

		terminating = false;
		notifyChange(true);
	}

	// Reset the program state, so that we may restart it later.
	void reset() {
		if (threads.any())
			terminate();

		// Clear output.
		output.clear();

		// Clear allocations.
		leakableAllocs.clear();

		// Reset the thread ID:s
		usedThreadId = 0;

		// Restore globals.
		if (globals)
			globals.restore(findGlobals);
	}

	// Get any allocations that might be leaks since the last call to this function.
	unsafe:RawPtr[] leakableAllocations() {
		if (leakableAllocs.any) {
			var old = leakableAllocs;
			leakableAllocs = unsafe:RawPtr[];
			return old;
		} else {
			return leakableAllocs;
		}
	}

	// This function will be called by the executed code to notify us of their current state in
	// their execution. This also lets us execute them at any location. 'offset' is used to decide
	// if it is necessary to break a barrier step operation.
	package void newLocation(SrcPosWrap pos, StackVar[] vars, Nat offset) {
		Word id = currentUThread();
		ProgThread t = threads.get(id);

		t.onNewLocation(pos.v, vars, offset);
	}

	// This function will be called by the executed code to notify that a (traced) function hit a
	// barrier.
	package void newBarrier(SrcPosWrap pos, Barrier type, Nat offset) {
		Word id = currentUThread();
		ProgThread t = threads.get(id);

		t.onNewBarrier(pos.v, type, offset);
	}

	// This function will be called by the executed code to notify that a (traced) function is about
	// to call a function marked as a barrier. Similar to above, but provides parameters as well.
	package void newFnBarrier(SrcPosWrap pos, Str name, Barrier type, StackVar[] params, Nat offset) {
		Word id = currentUThread();
		ProgThread t = threads.get(id);

		t.onNewFnBarrier(pos.v, name, type, params, offset);
	}

	// This function will be called by the executed code to notify that it is about to call a
	// function. The goal is only to inform the system about the progress inside the function so
	// that we can differentiate between different calls to the same function from the same
	// statement.
	package void beforeCall(Nat offset) {
		Word id = currentUThread();
		ProgThread t = threads.get(id);

		t.beforeCall(offset);
	}

	// This function will be called when the function is about to return. It is basically the same
	// as 'newLocation', except that we get the return value here, and we might want to indicate
	// that the function is returning somehow.
	package void functionReturn(StackVar[] vars) {
		Word id = currentUThread();
		ProgThread t = threads.get(id);

		t.onFunctionReturn(vars);
	}

	// This function will be called by the executed code to notify that a new (traced) function was
	// called, and that we shall open a new scope.
	package MemTracker functionEntered(Str name, SrcPosWrap pos, progvis:data:ViewHints hints, StackVar[] vars) {
		Word id = currentUThread();
		ProgThread t = if (t = threads.at(id)) {
			t;
		} else {
			ProgThread t(this, ++usedThreadId);
			addThread(id, t);
			t;
		};

		t.onFunctionEntered(name, pos.v, hints, vars);
	}

	// This function will be called on function entry of a non-traced function. 'functionExited'
	// will still not be called for this function so that we can properly track when the thread
	// exists.
	package MemTracker anonFunctionEntered() {
		Word id = currentUThread();
		ProgThread t = if (t = threads.at(id)) {
			t;
		} else {
			ProgThread t(this, ++usedThreadId);
			addThread(id, t);
			t;
		};

		t.onAnonFunctionEntered();
	}

	// This function will be called by the executed code to notify that a (traced) function is about
	// to exit.
	package void functionExited() {
		Word id = currentUThread();
		unless (t = threads.at(id))
			return;

		var notify = t.onFunctionExited();
		if (!t.alive) {
			removeThread(id, t);

			if (f = t.watch)
				spawn checkErrors(t.threadId, f);

			lastTerminatedId = t.threadId;
			if (notify)
				notifyChange(true);
		}
	}

	// This function will be called by the executed code to notify that an anonymous function is
	// about to exit.
	package void anonFunctionExited() {
		Word id = currentUThread();
		unless (t = threads.at(id))
			return;

		var notify = t.onAnonFunctionExited();
		if (!t.alive) {
			removeThread(id, t);

			if (f = t.watch)
				spawn checkErrors(t.threadId, f);

			lastTerminatedId = t.threadId;
			if (notify)
				notifyChange(true);
		}
	}

	// Allocate a new thread ID for a thread that will eventually start.
	Nat allocFutureId(Word threadId) {
		// We can just preemptively add the thread.
		Nat id = ++usedThreadId;
		ProgThread t(this, id);
		addThread(threadId, t);

		id;
	}

	// This function is called when the code performs an allocation that risks becoming a leak at
	// some point. This exists so that short-lived allocations will still be tracked properly if
	// their whole existence is stepped over.
	package void leakableAllocation(unsafe:RawPtr alloc) {
		leakableAllocs << alloc;
	}

	// Check for errors from a future. Will block until anything is posted.
	private void checkErrors(Nat threadId, FutureBase from) {
		try {
			from.errorResult();
		} catch (ExitError e) {
			// Just ignore it.
		} catch (Exception e) {
			notifyProgramError(lastTerminatedId, e);
		}
	}
}

/**
 * Wrapper to make a pointer of an SrcPos for easier management from ASM.
 */
class SrcPosWrap {
	SrcPos v;

	init(SrcPos pos) {
		init() { v = pos; }
	}
}


package PkgReader? progvisReader(Str ext, Url[] files, Package pkg) on Compiler {
	unless (name = readerName(ext))
		return null;

	name.add(0, "progvis");
	return createReader(name, files, pkg);
}

// Find a standard Storm reader.
package PkgReader? stdReader(Str ext, Url[] files, Package pkg) on Compiler {
	if (name = readerName(ext))
		return createReader(name, files, pkg);
	null;
}

private void addProgvisHints(Str ext, Hints[] to) on Compiler {
	unless (name = readerName(ext))
		return;

	name.add(0, "progvis");
	name.last.name = "hints";
	name.last.params.clear();
	unless (fn = rootScope.find(name) as Function)
		return;

	unless (ptr = fn.pointer as Fn<Hints>)
		return;

	to << ptr.call();
}

private void addStdHints(Str ext, Hints[] to) on Compiler {
	unless (name = readerName(ext))
		return;

	name.last.name = "hints";
	name.last.params.clear();
	unless (fn = rootScope.find(name) as Function)
		return;

	unless (ptr = fn.pointer as Fn<Hints>)
		return;

	to << ptr.call();
}

// Helper to load a program on a different thread.
private Program loadProgram(Url[] files) on Compiler {
	Package pkg("<sandbox>");
	pkg.parentLookup = rootPkg;
	pkg.retainSource();

	Map<Str, Url[]> exts;
	for (f in files) {
		exts[f.ext] << f;
	}

	Hints[] hints;
	PkgReader[] readers;
	for (ext, files in exts) {
		// Try to find a reader ourselves first.
		if (r = progvisReader(ext, files, pkg)) {
			readers << r;
			addProgvisHints(ext, hints);
		} else if (r = stdReader(ext, files, pkg)) {
			readers << r;
			addStdHints(ext, hints);
		} else {
			throw LoadError("The file type ${ext} is not supported.");
		}
	}

	if (readers.empty()) {
		throw LoadError("The desired file type(s) are not supported.");
	}

	// Add the standard hints last, so that they have a chance to do things.
	hints << defaultHints();

	// Load all code.
	read(readers);

	// Produce errors now rather than later.
	pkg.compile();

	// Extract the view hints for later use.
	progvis:data:ViewHints[] viewHints;
	for (h in hints)
		if (view = h.view)
			viewHints << view;

	// Create the program instance and patch all code.
	Program program(pkg, findMain(hints, pkg), viewHints);
	patchFunctions(program, hints, pkg);
	program.loadSource(files);
	return program;
}

// Find the main function.
private Fn<void> findMain(Hints[] hints, Package pkg) {
	for (h in hints) {
		if (r = h.code.findMain(pkg))
			return r;
	}

	throw progvis:ProgvisError("Unable to find a 'main' function!");
}


// Find supported file types. On compiler thread.
private Set<Str> findSupportedFiles() on Compiler {
	Package pkg = named{lang};
	Set<Str> to;
	findSupportedFiles(to, named{:lang});
	findSupportedFiles(to, named{:progvis:lang});
	to;
}

private void findSupportedFiles(Set<Str> out, Package pkg) on Compiler {
	pkg.loadAll;
	for (elem in pkg) {
		if (elem as Package)
			out.put(elem.name);
	}
}

// Find global variables in a package.
private GlobalVar[] findGlobalsHelp(Package pkg) on Compiler {
	GlobalVar[] result;

	for (elem in pkg) {
		if (elem as GlobalVar) {
			if (elem.owner is named{Render})
				result << elem;
		}
	}

	result;
}

/**
 * Exception thrown during load.
 */
class LoadError extends Exception {
	init(Str msg) { init { msg = msg; } }
	Str msg;
	void message(StrBuf to) : override {
		to << msg;
	}
}


// Global variable that keeps track of all Program instances, so that we can find the owner to a
// particular thread by its thread id.
// TODO: This can't be private...
package WeakSet<Program> programInstances on Render = WeakSet<Program>();

// Find the Program instance containing a particular thread ID.
private Program? findProgram(Word threadId) on Render {
	for (i in programInstances) {
		if (i.threads.has(threadId))
			return i;
	}
	return null;
}

public Program? findProgram() {
	findProgram(currentUThread());
}

// Find a particular thread.
public ProgThread? findThread(Word threadId) on Render {
	for (i in programInstances) {
		if (t = i.threads.at(threadId))
			return t;
	}
	return null;
}

private Nat findThisThreadId(Word id) on Render {
	if (t = findThread(id)) {
		return t.threadId;
	}

	0;
}

// Find a thread id of this (existing) thread.
public Nat findThisThreadId() {
	findThisThreadId(currentUThread());
}

// Find the "simple" thread ID for a thread. This might create the thread ID in the appropriate
// location, which is why the current thread ID is needed as well.
private Nat findNewThreadId(Word currentThread, Word query) on Render {
	if (thread = findThread(query)) {
		return thread.threadId;
	}

	// We need to create it.
	if (p = findProgram(currentThread)) {
		return p.allocFutureId(query);
	}

	0;
}

// Wrapper not bound to a particular thread, making it easier to call from ASM.
public Nat findNewThreadId(Word query) {
	findNewThreadId(currentUThread(), query);
}

private void watchErrorsI(Word threadId, FutureBase future) on Render {
	if (thread = findThread(threadId)) {
		thread.watch = future;
	} else {
		future.detach();
	}
}

// Tell the system to watch for errors for the specified thread ID using the supplied future.
public void watchErrors(Word threadId, FutureBase future) {
	watchErrorsI(threadId, future);
}

// Tell the system that we are about to throw some kind of exception, and that this will likely terminate the program.
public void onFatalException() {
	if (p = findThread(currentUThread()))
		p.onFatalException();
}

// Add output to the program associated with the current thread, or output to stdout if no thread
// currently available. No newlines are added automatically.
public void output(Str text) {
	if (program = findProgram(currentUThread())) {
		program.addOutput(text);
	} else {
		stdOut.write(text);
		stdOut.flush();
	}
}