import java.lang.invoke.MethodHandle;
import java.lang.invoke.MethodHandles;
import jdk.internal.vm.Continuation;
import jdk.internal.vm.ContinuationScope;
import util.UnsafeUtil;

// java --add-opens java.base/jdk.internal.vm=ALL-UNNAMED Loom3
public class Loom3 {
	private static final ContinuationScope cs = new ContinuationScope("VirtualThreads");
	private static final long CONT_OFFSET;
	// private static final long MOUNTED_OFFSET;
//	private static final MethodHandle DO_YIELD;
	private static final MethodHandle YIELD0;
	private static final MethodHandle ENTER_SPECIAL;
//	private static final MethodHandle SET_EXTENT_LOCAL_CACHE;

	static {
		try {
			CONT_OFFSET = UnsafeUtil.objectFieldOffset(Thread.class, "cont");
//			MOUNTED_OFFSET = U.objectFieldOffset(UnsafeUtil.getField(Continuation.class, "mounted"));
			var lookup = MethodHandles.lookup();
//			DO_YIELD = lookup.unreflect(UnsafeUtil.getMethod(Continuation.class, "doYield"));
			YIELD0 = lookup.unreflect(UnsafeUtil.getMethod(Continuation.class, "yield0",
					ContinuationScope.class, Continuation.class));
			ENTER_SPECIAL = lookup.unreflect(UnsafeUtil.getMethod(Continuation.class, "enterSpecial",
					Continuation.class, boolean.class, boolean.class));
//			m = UnsafeUtil.getMethod(Thread.class, "setExtentLocalCache", Object[].class);
//			SET_EXTENT_LOCAL_CACHE = MethodHandles.lookup().unreflect(m);
		} catch (ReflectiveOperationException e) {
			throw new RuntimeException(e);
		}
	}

	static Continuation start(Runnable r) {
		var c = new Continuation(cs, r);
		c.run();
//		try {
//			U.putBooleanVolatile(c, MOUNTED_OFFSET, true);
//			SET_EXTENT_LOCAL_CACHE.invokeExact((Object[])null);
//			U.putObject(Thread.currentThread(), CONT_OFFSET, c);
//			ENTER_SPECIAL.invokeExact(c, false, false);
//		} catch (Throwable e) {
//			throw new RuntimeException(e);
//		}
		return c;
	}

	@SuppressWarnings("removal")
	static void pause() {
		var c = (Continuation)UnsafeUtil.unsafe.getObject(Thread.currentThread(), CONT_OFFSET);
		try {
			@SuppressWarnings("unused")
			var __ = (boolean)YIELD0.invokeExact(c, cs, (Continuation)null);
			// var __ = (int)DO_YIELD.invokeExact();
			// VarHandle.releaseFence(); // needed to prevent certain transformations by the compiler
		} catch (Throwable e) {
			throw new RuntimeException(e);
		} finally {
			UnsafeUtil.unsafe.putObject(Thread.currentThread(), CONT_OFFSET, c);
		}
	}

	static void resume(Continuation c) {
		if (c.isDone())
			throw new IllegalStateException("Continuation terminated");
		try {
			ENTER_SPECIAL.invokeExact(c, true, false);
		} catch (Throwable e) {
			throw new RuntimeException(e);
		}
	}

	@SuppressWarnings("removal")
	private static void test(int i) {
		System.out.println("begin " + i);
		var c1 = start(() -> {
			System.out.println("before pause1: " + UnsafeUtil.unsafe.getObject(Thread.currentThread(), CONT_OFFSET));
			pause();
			System.out.println("after pause1: " + UnsafeUtil.unsafe.getObject(Thread.currentThread(), CONT_OFFSET));
		});
		var c2 = start(() -> {
			System.out.println("before pause2: " + UnsafeUtil.unsafe.getObject(Thread.currentThread(), CONT_OFFSET));
			pause();
			System.out.println("after pause2: " + UnsafeUtil.unsafe.getObject(Thread.currentThread(), CONT_OFFSET));
		});
		System.out.println("resume1");
		resume(c1);
		System.out.println("resume2");
		resume(c2);
		System.out.println("end");
	}

	public static void main(String[] args) {
//		for (int i = 0; i < 1000; i++)
		test(1);
	}
}