package eu.jonahbauer.wizard.core.machine;

import eu.jonahbauer.wizard.core.machine.states.State;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import org.jetbrains.annotations.Blocking;
import org.jetbrains.annotations.NonBlocking;
import org.jetbrains.annotations.NotNull;

import java.util.Comparator;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;

public abstract class Context<S extends State<S,C>, C extends Context<S,C>> {
    private final ScheduledExecutorService scheduler = Executors.newSingleThreadScheduledExecutor();
    private final ThreadPoolExecutor executor = new ThreadPoolExecutor(
            1, 1,
            0, TimeUnit.SECONDS,
            new PriorityBlockingQueue<>(
                    11,
                    Comparator.comparingInt(r -> r instanceof PriorityRunnable prio ? prio.getPriority() : Integer.MAX_VALUE)
                              .thenComparingLong(r -> r instanceof PriorityRunnable prio ? prio.getTimestamp() : Long.MAX_VALUE)
            ),
            r -> {
                var t = new Thread(r);
                t.setUncaughtExceptionHandler((t1, e) -> finish(e));
                return t;
            }
    );

    protected S state;

    private final CompletableFuture<Void> future = new CompletableFuture<>();
    private final CompletableFuture<Void> finished = future.whenComplete((v, t) -> {
        executor.shutdownNow();
        scheduler.shutdownNow();
    });

    @NonBlocking
    protected CompletableFuture<Void> submit(Runnable runnable) {
        var future = new CompletableFuture<Void>();
        executor.execute(new PriorityRunnable(100, () -> {
            try {
                runnable.run();
                future.complete(null);
            } catch (Throwable t) {
                future.completeExceptionally(t);
            }
        }));
        return future;
    }

    @Blocking
    protected void start(@NotNull S state) {
        CountDownLatch latch = new CountDownLatch(1);
        AtomicReference<RuntimeException> exception = new AtomicReference<>();

        executor.execute(new PriorityRunnable(0, () -> {
            if (future.isDone()) {
                exception.set(new IllegalStateException("Context has already finished."));
                latch.countDown();
            } else {
                latch.countDown();
                doTransition(null, state);
            }
        }));

        while (true) {
            try {
                latch.await();
                if (exception.get() != null) {
                    throw exception.get();
                }
                break;
            } catch (InterruptedException ignored) {}
        }
    }

    @NonBlocking
    public void transition(S currentState, S newState) {
        executor.execute(new PriorityRunnable(0, () -> doTransition(currentState, newState)));
    }

    @NonBlocking
    public void finish() {
        finish(null);
    }

    @NonBlocking
    public void finish(Throwable exception) {
        executor.execute(new PriorityRunnable(0, () -> doFinish(exception)));
    }

    @NonBlocking
    public void cancel() {
        finish(new CancellationException());
    }

    /*
     * internal methods that are called on the executor
     */

    private void doTransition(S currentState, S newState) {
        if (state == currentState) {
            state = newState;
            if (currentState != null) //noinspection unchecked
                currentState.onExit((C) this);
            onTransition(currentState, newState);
            if (newState != null) //noinspection unchecked
                newState.onEnter((C) this);
        }
    }

    private void doFinish(Throwable t) {
        if (future.isDone()) return;

        doTransition(state, null);
        if (t != null) {
            future.completeExceptionally(t);
        } else {
            future.complete(null);
        }
    }

    protected void onTransition(S from, S to) {}

    @Blocking
    public void await() throws InterruptedException, ExecutionException, CancellationException {
        finished.get();
    }

    public void timeout(@NotNull S currentState, long delay) {
        scheduler.schedule(() -> {
            submit(() -> {
                if (state == currentState) {
                    //noinspection unchecked
                    state.onTimeout((C) this);
                }
            });
        }, delay, TimeUnit.MILLISECONDS);
    }

    @Getter
    @RequiredArgsConstructor
    private static class PriorityRunnable implements Runnable {
        private final int priority;
        private final long timestamp = System.nanoTime();
        private final Runnable runnable;

        @Override
        public void run() {
            runnable.run();
        }
    }
}
