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) {
        Set<State> allStates = new java.util.HashSet<>();
        if (transitions != null) {
            transitions.forEach(t -> {
                if (t.getSourceStates() != null) allStates.addAll(t.getSourceStates());
                if (t.getTargetStates() != null) allStates.addAll(t.getTargetStates());
            });
        }
        if (initialStates != null) {
            initialStates.forEach(s -> allStates.add(State.of(s)));
        }
        if (endStates != null) {
            endStates.forEach(s -> allStates.add(State.of(s)));
        }
        return allStates;
    }

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