package pl.polgrabia.demos.crackingcodeinterview.t16x20;

import java.util.function.Consumer;

public class RstTree {
    private RstTreeNode rstTreeRoot;

    public void add(String word) {
        assert word != null : "word must not be null";
        assert !word.isBlank() : "word must not be blank";
        assert word.matches("^[a-z]+$") : "word '" + word + "' must contain only small letters";
        if (rstTreeRoot == null) {
            rstTreeRoot = new RstTreeNode();
            rstTreeRoot.setPrefix("");
            rstTreeRoot.setValue(word);
        } else {
            addInternal(rstTreeRoot, word);
        }
    }

    private static boolean addInternal(RstTreeNode node, String word) {
        if (word.equals(node.getValue())) {
            // already added
            return false;
        }

        assert word.startsWith(node.getPrefix()) : "word must start with prefix";
        int prefixLength = node.getPrefix().length();
        if (prefixLength >= word.length()) {
            String oldValue = node.getValue();
            node.setValue(word);
            boolean added = addInternal(node, oldValue);
            if (!added) {
                node.setValue(oldValue);
                return false;
            }

            return true;
        }
        char c = word.charAt(prefixLength);
        int offset = c - 'a';
        assert offset >= 0 && offset < RstTreeNode.EN_ALPHABET_LETTER_COUNT : "It's an offset for possible prefix child item - max alphabet";
        // logger.info("Word: {}, prefix: {}, index: {}, c: {}, offset: {}", word, node.getPrefix(), node.getPrefix().length(), c, offset);
        var childItem = node.getRstChildren().length <= offset ? null : node.getRstChildren()[offset];
        if (childItem == null) {
            childItem = new RstTreeNode();
            childItem.setPrefix(word.substring(0, prefixLength + 1));
            childItem.setValue(word);
            node.getRstChildren()[offset] = childItem;
            return true;
        } else {
            return addInternal(childItem, word);
        }
    }

    public boolean contains(String word) {
        assert word != null : "word must not be null";
        assert !word.isBlank() : "word must not be blank";
        assert word.matches("^[a-z]+$") : "word '" + word + "' must contain only small letters";

        RstTreeNode searchedNode = lookupNodesBySearchKeyInternal(rstTreeRoot, word, SearchMode.VALUE);
        if (searchedNode == null) {
            return false;
        }

        String value = searchedNode.getValue();
        return word.equals(value);
    }

    private RstTreeNode lookupNodesBySearchKeyInternal(RstTreeNode node, String searchValue, SearchMode searchMode) {
        if (node == null) {
            return null;
        }

        // logger.info("Node value: {}, prefix: {}, searchValue: {}", node.getValue(), node.getPrefix(), searchValue);

        assert searchValue.startsWith(node.getPrefix()) : "searchValue " + searchValue + "must start with prefix: " + node.getPrefix();

        if (SearchMode.VALUE.equals(searchMode) && searchValue.equals(node.getValue())) {
            return node;
        }

        if (SearchMode.PREFIX_OR_VALUE.equals(searchMode)
                && (node.getPrefix().startsWith(searchValue) || node.getValue().startsWith(searchValue))) {
            return node;
        }

        if (SearchMode.PREFIX.equals(searchMode)
                && (node.getPrefix().startsWith(searchValue))) {
            return node;
        }

        if (node.getPrefix().startsWith(searchValue)) {
            return null;
        }

        char c = searchValue.charAt(node.getPrefix().length());
        int offset = c - 'a';
        assert offset >= 0 && offset < RstTreeNode.EN_ALPHABET_LETTER_COUNT : "It's an offset for possible prefix child item - max alphabet";
        // logger.info("Word: {}, prefix: {}, index: {}, c: {}, offset: {}", searchValue, node.getPrefix(), node.getPrefix().length(), c, offset);
        var childItem = node.getRstChildren().length <= offset ? null : node.getRstChildren()[offset];
        // logger.info("Childitem: {}, c: {}, offset: {}", childItem, c, offset);
        return lookupNodesBySearchKeyInternal(childItem, searchValue, SearchMode.VALUE);
    }

    public void dfsByPrefix(String prefix, Consumer<String> wordConsumer) {
        assert prefix != null : "prefix must not be null";
        assert !prefix.isBlank() : "prefix must not be blank";
        assert prefix.matches("^[a-z]+$") : "prefix '" + prefix + "' must contain only small letters";

        RstTreeNode node = lookupNodesBySearchKeyInternal(rstTreeRoot, prefix, SearchMode.PREFIX_OR_VALUE);
        dfs(node, wordConsumer);
    }

    private void dfs(RstTreeNode node, Consumer<String> wordConsumer) {
        if (node == null) {
            return;
        }
        wordConsumer.accept(node.getValue());
        for (int i = 0; i < node.getRstChildren().length; i++) {
            dfs(node.getRstChildren()[i], wordConsumer);
        }
    }

    private enum SearchMode {
        VALUE,
        PREFIX_OR_VALUE,
        PREFIX
    }
}
