package click.kamil.springstatemachineexporter.ast.app;

import click.kamil.springstatemachineexporter.model.State;
import click.kamil.springstatemachineexporter.model.Transition;

import java.util.Collection;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;

public class TransitionStateUtils {
    public static Set<String> findStartStates(Collection<Transition> transitions, Set<String> initialStates) {
        if (initialStates != null && !initialStates.isEmpty()) {
            return initialStates;
        }
        return findStartStates(transitions);
    }

    public static Set<String> findStartStates(Collection<Transition> transitions) {
        Set<String> allTargetStates = transitions.stream()
                .flatMap(t -> normalizeStates(t.getTargetStates()))
                .collect(Collectors.toSet());
        return transitions.stream()
                .flatMap(t -> normalizeStates(t.getSourceStates()))
                .filter(state -> !allTargetStates.contains(state))
                .collect(Collectors.toSet());
    }

    public static Set<String> findEndStates(Collection<Transition> transitions, Set<String> endStates) {
        if (endStates != null && !endStates.isEmpty()) {
            return endStates;
        }
        return findEndStates(transitions);
    }

    public static Set<String> findEndStates(Collection<Transition> transitions) {
        Set<String> allSourceStates = transitions.stream()
                .flatMap(t -> normalizeStates(t.getSourceStates()))
                .collect(Collectors.toSet());
        return transitions.stream()
                .flatMap(t -> normalizeStates(t.getTargetStates()))
                .filter(state -> !allSourceStates.contains(state))
                .collect(Collectors.toSet());
    }

    public static Set<State> findAllStates(Collection<Transition> transitions, Set<String> initialStates, Set<String> endStates) {
        java.util.Map<String, State> byFullIdentifier = new java.util.LinkedHashMap<>();
        if (transitions != null) {
            transitions.forEach(t -> {
                mergeState(byFullIdentifier, t.getSourceStates());
                mergeState(byFullIdentifier, t.getTargetStates());
            });
        }
        if (initialStates != null) {
            initialStates.forEach(label -> mergeState(byFullIdentifier, State.of(label)));
        }
        if (endStates != null) {
            endStates.forEach(label -> mergeState(byFullIdentifier, State.of(label)));
        }
        return new java.util.LinkedHashSet<>(byFullIdentifier.values());
    }

    private static void mergeState(java.util.Map<String, State> map, State state) {
        if (state == null) {
            return;
        }
        String key = stateKey(state);
        map.putIfAbsent(key, state);
    }

    private static void mergeState(java.util.Map<String, State> map, java.util.Collection<State> states) {
        if (states == null) {
            return;
        }
        for (State state : states) {
            mergeState(map, state);
        }
    }

    private static String stateKey(State state) {
        if (state.fullIdentifier() != null && !state.fullIdentifier().isBlank()) {
            return state.fullIdentifier();
        }
        return state.rawName();
    }

    private static Stream<String> normalizeStates(Collection<State> states) {
        return states.stream()
                .map(State::toString)
                .map(QuotedExpression::stripQuotes)
                .filter(Objects::nonNull);
    }
}
