package org.luaj.vm2.libs.jse;

import java.util.Comparator;
import java.util.PriorityQueue;

import org.luaj.vm2.Globals;
import org.luaj.vm2.LuaError;
import org.luaj.vm2.LuaString;
import org.luaj.vm2.LuaThread;
import org.luaj.vm2.LuaValue;
import org.luaj.vm2.Varargs;
import org.luaj.vm2.libs.VarArgFunction;

/**
 * Cooperative host-side scheduler for Lua tasks that use explicit wait/yield calls.
 * <p>
 * This does not change luaj's coroutine implementation. It provides a small host
 * API for scheduling resumptions and timer-driven waits.
 */
public class LuaScheduler {

	public interface Clock {
		long nowMillis();
	}

	private static final Clock SYSTEM_CLOCK = new Clock() {
		public long nowMillis() {
			return System.currentTimeMillis();
		}
	};

	private static final LuaString WAIT_SIGNAL = LuaValue.valueOf("__luaj_scheduler_wait");
	private static final LuaString YIELD_SIGNAL = LuaValue.valueOf("__luaj_scheduler_yield");

	private final Globals globals;
	private final Clock clock;
	private final PriorityQueue<ScheduledTask> queue = new PriorityQueue<ScheduledTask>(11, new Comparator<ScheduledTask>() {
		public int compare(ScheduledTask a, ScheduledTask b) {
			if (a.wakeAtMillis != b.wakeAtMillis) {
				return a.wakeAtMillis < b.wakeAtMillis ? -1 : 1;
			}
			if (a.eligiblePass != b.eligiblePass) {
				return a.eligiblePass < b.eligiblePass ? -1 : 1;
			}
			if (a.sequence == b.sequence) {
				return 0;
			}
			return a.sequence < b.sequence ? -1 : 1;
		}
	});

	private long sequence;
	private long pass;

	public LuaScheduler(Globals globals) {
		this(globals, SYSTEM_CLOCK);
	}

	public LuaScheduler(Globals globals, Clock clock) {
		this.globals = globals;
		this.clock = clock;
		installIntoGlobals();
	}

	public Globals getGlobals() {
		return globals;
	}

	public ScheduledTask spawn(String script, String chunkname) {
		return spawn(globals.load(script, chunkname));
	}

	public ScheduledTask spawn(LuaValue function) {
		return spawn(function, LuaValue.NONE);
	}

	public ScheduledTask spawn(LuaValue function, Varargs args) {
		ScheduledTask task = new ScheduledTask(new LuaThread(globals, function.checkfunction()));
		task.pendingArgs = args;
		schedule(task, clock.nowMillis(), 0);
		return task;
	}

	public boolean hasPendingTasks() {
		return !queue.isEmpty();
	}

	public long millisUntilNextTask() {
		ScheduledTask next = queue.peek();
		if (next == null) {
			return -1;
		}
		long delay = next.wakeAtMillis - clock.nowMillis();
		return delay > 0 ? delay : 0;
	}

	public int runReadyTasks() {
		long now = clock.nowMillis();
		long currentPass = ++pass;
		int ran = 0;
		while (true) {
			ScheduledTask task = queue.peek();
			if (task == null || task.wakeAtMillis > now || task.eligiblePass > currentPass) {
				return ran;
			}
			queue.poll();
			if (task.done) {
				continue;
			}
			resumeTask(task, now, currentPass);
			ran++;
		}
	}

	private void resumeTask(ScheduledTask task, long now, long currentPass) {
		Varargs args = task.pendingArgs;
		task.pendingArgs = LuaValue.NONE;
		Varargs result = task.thread.resume(args);

		if (!result.arg1().toboolean()) {
			task.done = true;
			task.failure = new LuaError(result.arg(2).tojstring());
			return;
		}

		LuaValue signal = result.arg(2);
		if (signal.raweq(WAIT_SIGNAL)) {
			long delay = result.arg(3).optlong(0);
			schedule(task, now + Math.max(0, delay), currentPass + 1);
			return;
		}
		if (signal.raweq(YIELD_SIGNAL)) {
			task.lastYield = result.subargs(3);
			schedule(task, now, currentPass + 1);
			return;
		}

		if ("dead".equals(task.thread.getStatus())) {
			task.done = true;
			task.result = result.subargs(2);
			return;
		}

		task.lastYield = result.subargs(2);
		schedule(task, now, currentPass + 1);
	}

	private void schedule(ScheduledTask task, long wakeAtMillis, long eligiblePass) {
		task.wakeAtMillis = wakeAtMillis;
		task.eligiblePass = eligiblePass;
		task.sequence = sequence++;
		queue.add(task);
	}

	private void installIntoGlobals() {
		globals.set("wait", new WaitFunction());
		globals.set("yield", new YieldFunction());
		globals.set("spawn", new SpawnFunction());
	}

	private final class WaitFunction extends VarArgFunction {
		public Varargs invoke(Varargs args) {
			long delay = args.arg(1).optlong(0);
			return globals.yield(LuaValue.varargsOf(WAIT_SIGNAL, LuaValue.valueOf(Math.max(0, delay))));
		}
	}

	private final class YieldFunction extends VarArgFunction {
		public Varargs invoke(Varargs args) {
			return globals.yield(LuaValue.varargsOf(YIELD_SIGNAL, args));
		}
	}

	private final class SpawnFunction extends VarArgFunction {
		public Varargs invoke(Varargs args) {
			ScheduledTask task = spawn(args.checkfunction(1), args.subargs(2));
			return task.thread;
		}
	}

	public static final class ScheduledTask {
		private final LuaThread thread;
		private Varargs pendingArgs = LuaValue.NONE;
		private Varargs result = LuaValue.NONE;
		private Varargs lastYield = LuaValue.NONE;
		private LuaError failure;
		private boolean done;
		private long wakeAtMillis;
		private long eligiblePass;
		private long sequence;

		private ScheduledTask(LuaThread thread) {
			this.thread = thread;
		}

		public LuaThread getThread() {
			return thread;
		}

		public boolean isDone() {
			return done;
		}

		public boolean isFailed() {
			return failure != null;
		}

		public LuaError getFailure() {
			return failure;
		}

		public Varargs getResult() {
			return result;
		}

		public Varargs getLastYield() {
			return lastYield;
		}

		public long getWakeAtMillis() {
			return wakeAtMillis;
		}
	}
}
