package eu.jonahbauer.wizard.common.machine;

import lombok.Getter;
import org.jetbrains.annotations.NotNull;

import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.locks.ReentrantLock;
import java.util.function.Function;
import java.util.function.Supplier;

public abstract class Context<S extends State<S,C>, C extends Context<S,C>> {
    @Getter
    private @NotNull S state;
    private final ReentrantLock lock = new ReentrantLock();

    public Context(@NotNull S initialState) {
        this.state = initialState;
    }

    /**
     * Atomically executes the given supplier and transitions to the returned state if present.
     * @see #transition(State)
     */
    public void execute(Supplier<Optional<@NotNull S>> transition) {
        lock.lock();
        try {
            Optional<S> next = transition.get();
            next.ifPresent(this::transition);
        } finally {
            lock.unlock();
        }
    }

    /**
     * Atomically applies the given function to the current state and transitions to the returned state if present.
     * @see #transition(State)
     */
    public void execute(Function<@NotNull S, Optional<@NotNull S>> transition) {
        lock.lock();
        try {
            Optional<S> next = transition.apply(state);
            next.ifPresent(this::transition);
        } finally {
            lock.unlock();
        }
    }

    /**
     * Atomically transitions to the specified state calling the state lifecycle methods.
     * When an error occurs during execution of the lifecycle methods {@link #handleError(Throwable)} is called.
     * @param state the next state
     * @see State#onExit(Context)
     * @see State#onEnter(Context)
     * @see #handleError(Throwable)
     */
    @SuppressWarnings("unchecked")
    public void transition(@NotNull S state) {
        lock.lock();
        try {
            this.state.onExit((C) this);
            onTransition(this.state, state);
            this.state = state;
            state.onEnter((C) this).ifPresent(this::transition);
        } catch (Throwable t) {
            handleError(t);
        } finally {
            lock.unlock();
        }
    }

    /**
     * Atomically checks that the current state is {@linkplain Object#equals(Object) equal to} the specified expected
     * state and transitions to the specified next state.
     * @param expected the expected current state
     * @param next the next state
     * @see #transition(State)
     * @throws IllegalStateException if the current state is not equal to the expected state
     */
    public void transition(@NotNull S expected, @NotNull S next) {
        lock.lock();
        try {
            if (Objects.equals(expected, this.state)) {
                transition(next);
            } else {
                throw new IllegalStateException("Current state is not " + expected + ".");
            }
        } finally {
            lock.unlock();
        }
    }

    /**
     * Atomically transitions to the specified state without calling the state lifecycle methods. This method can
     * be used to recover from errors occurring during {@link State#onExit(Context)} which would otherwise prevent
     * the context from transitioning to another state.
     * @param state the next state
     */
    protected void forceTransition(@NotNull S state) {
        lock.lock();
        try {
            onTransition(this.state, state);
            this.state = state;
        } finally {
            lock.unlock();
        }
    }

    /**
     * Callback method that will synchronously be called on transitioning between two states.
     * @param from the previous state
     * @param to the next state
     */
    protected void onTransition(S from, S to) {}

    /**
     * Callback method that will synchronously be called when an error occurs during execution of state lifecycle
     * methods.
     * @param t the cause
     */
    protected abstract void handleError(Throwable t);
}
