package gudusoft.gsqlparser.lineage.v2.json;

import gudusoft.gsqlparser.lineage.v2.SqllrResources;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.SortedMap;
import java.util.SortedSet;
import java.util.TreeMap;
import java.util.TreeSet;

/**
 * The array registry of {@code canonical-json-v1}, Section 6: the class ({@code SET}, {@code SEQUENCE},
 * {@code DISPLAY_ORDER}) and the sort key of every array position of the lineage document, read from the
 * {@code x-sqllr-order} and {@code x-sqllr-sortKey} annotations of the schema, plus the positions of open
 * content (Section 5, rule 5), where every array is a {@code SEQUENCE}.
 *
 * <p>The positions form a tree: {@link #root()} is the document, {@link Position#member(String)} a member
 * and {@link Position#element()} any element of an array. A position outside the tree is unregistered.</p>
 *
 * <p>The derivation is the reference's {@code schema_array_registry}; the open-content positions are its
 * {@code OPEN_CONTENT}, which the schema does not mark.</p>
 *
 * @since 4.2.17
 */
public final class ArrayRegistry {

    /** The registry group {@code arrayClasses}. */
    public enum Order {
        SET, SEQUENCE, DISPLAY_ORDER
    }

    /** The sort key of a {@code SET} (Section 6.2). */
    public static final class SortKey {
        /** What the key of an element is. */
        public enum Kind {
            /** The element itself. */
            SELF,
            /** The element's {@code ref} member. */
            REF,
            /** Nothing beyond the canonical bytes of the element. */
            BYTES,
            /** A tuple of members of the element. */
            MEMBERS
        }

        private final Kind kind;
        private final List<String> members;

        private SortKey(Kind kind, List<String> members) {
            this.kind = kind;
            this.members = members;
        }

        public Kind getKind() {
            return kind;
        }

        /** The member names of a {@link Kind#MEMBERS} key; empty otherwise. */
        public List<String> getMembers() {
            return members;
        }

        static SortKey of(Object annotation, String where) {
            if ("self".equals(annotation)) {
                return new SortKey(Kind.SELF, Collections.<String>emptyList());
            }
            if ("ref".equals(annotation)) {
                return new SortKey(Kind.REF, Collections.<String>emptyList());
            }
            if ("bytes".equals(annotation)) {
                return new SortKey(Kind.BYTES, Collections.<String>emptyList());
            }
            if (annotation instanceof List<?> && !((List<?>) annotation).isEmpty()) {
                List<String> names = new ArrayList<String>();
                for (Object name : (List<?>) annotation) {
                    if (!(name instanceof String)) {
                        throw new IllegalStateException("x-sqllr-sortKey of " + where + " is not a list of names");
                    }
                    names.add((String) name);
                }
                return new SortKey(Kind.MEMBERS, Collections.unmodifiableList(names));
            }
            throw new IllegalStateException("unknown x-sqllr-sortKey " + annotation + " at " + where);
        }

        @Override
        public boolean equals(Object other) {
            return other instanceof SortKey && ((SortKey) other).kind == kind && ((SortKey) other).members.equals(members);
        }

        @Override
        public int hashCode() {
            return kind.hashCode() * 31 + members.hashCode();
        }

        /** {@code self}, {@code ref}, {@code bytes}, or {@code (a, b)}. */
        @Override
        public String toString() {
            switch (kind) {
                case SELF: return "self";
                case REF: return "ref";
                case BYTES: return "bytes";
                default:
                    StringBuilder out = new StringBuilder("(");
                    for (int i = 0; i < members.size(); i++) {
                        out.append(i == 0 ? "" : ", ").append(members.get(i));
                    }
                    return out.append(')').toString();
            }
        }
    }

    /** The registration of one array position. */
    public static final class Registration {
        /** Every array in open content. */
        public static final Registration OPEN_SEQUENCE = new Registration(Order.SEQUENCE, null);

        private final Order order;
        private final SortKey sortKey;

        Registration(Order order, SortKey sortKey) {
            this.order = order;
            this.sortKey = sortKey;
        }

        public Order getOrder() {
            return order;
        }

        /** The sort key of a {@code SET}; null for the other classes. */
        public SortKey getSortKey() {
            return sortKey;
        }

        @Override
        public boolean equals(Object other) {
            if (!(other instanceof Registration)) {
                return false;
            }
            Registration that = (Registration) other;
            return order == that.order && (sortKey == null ? that.sortKey == null : sortKey.equals(that.sortKey));
        }

        @Override
        public int hashCode() {
            return order.hashCode() * 31 + (sortKey == null ? 0 : sortKey.hashCode());
        }

        /** {@code SET ref}, {@code SEQUENCE}. */
        @Override
        public String toString() {
            return sortKey == null ? order.name() : order.name() + " " + sortKey;
        }
    }

    /** A position of the document: the root, a member of a position, or any element of an array position. */
    public static final class Position {
        private final Map<String, Position> members = new LinkedHashMap<String, Position>();
        private Position element;
        private Registration registration;
        private boolean openContent;

        private Position() {
        }

        /** The position of member {@code name}, or null when nothing is registered at or below it. */
        public Position member(String name) {
            return members.get(name);
        }

        /** The position of any element of this array, or null when nothing is registered at or below it. */
        public Position element() {
            return element;
        }

        /** The registration of the array at this position, or null. */
        public Registration registration() {
            return registration;
        }

        /** True when the value at this position, and everything inside it, is open content. */
        public boolean isOpenContent() {
            return openContent;
        }

        private Position child(String step) {
            if (ELEMENT.equals(step)) {
                if (element == null) {
                    element = new Position();
                }
                return element;
            }
            Position child = members.get(step);
            if (child == null) {
                child = new Position();
                members.put(step, child);
            }
            return child;
        }
    }

    /** The step for "any element" in the path constants of this class. */
    private static final String ELEMENT = "[]";

    /** The module collections whose element schemas are not defined yet (canonical-json-v1, Section 6.3). */
    public static final List<String> MODULE_ARRAYS = Collections.unmodifiableList(Arrays.asList(
            "dynamicSqlSites", "evidence", "expressionSpans", "naryFacts", "refinements", "routineCalls",
            "structuralSignatures", "typedSteps", "wildcardSelections"));

    /** Open content other than {@code extensions} and the elements of {@link #MODULE_ARRAYS}. */
    private static final String[][] OPEN_CONTENT = {
            {"merge"},
            {"reproducibilityManifest"},
            {"sourceUnits", ELEMENT, "originMap", ELEMENT},
            {"endpoints", ELEMENT, "pathSegments", ELEMENT},
            {"factCandidates", ELEMENT, "source", "candidateSet", "details"},
            {"factCandidates", ELEMENT, "target", "candidateSet", "details"},
            {"explainGraph", "candidateSets", ELEMENT, "details"},
            {"candidateSets", ELEMENT, "details"},
    };

    private final Position root;
    private final SortedMap<String, Registration> registrations;
    private final SortedSet<String> openContent;

    private ArrayRegistry(Position root, SortedMap<String, Registration> registrations, SortedSet<String> openContent) {
        this.root = root;
        this.registrations = Collections.unmodifiableSortedMap(registrations);
        this.openContent = Collections.unmodifiableSortedSet(openContent);
    }

    /** The position of the document root. */
    public Position root() {
        return root;
    }

    /** Every registered array position, as a display path ({@code $.endpoints[].name.segments}). */
    public SortedMap<String, Registration> registrations() {
        return registrations;
    }

    /** Every open-content position other than {@code extensions} members, as a display path. */
    public SortedSet<String> openContent() {
        return openContent;
    }

    /** The registry of the schema shipped with this library. */
    public static ArrayRegistry standard() {
        return Standard.INSTANCE;
    }

    private static final class Standard {
        static final ArrayRegistry INSTANCE = fromSchema(SqllrResources.lineageDocumentSchema());
    }

    /**
     * Derives the registry from a lineage document schema (a parsed JSON value).
     *
     * @throws IllegalStateException when an array position has no class, two classes, or an unknown sort key
     */
    public static ArrayRegistry fromSchema(Object schema) {
        Map<?, ?> top = (Map<?, ?>) schema;
        Map<?, ?> defs = (Map<?, ?>) top.get("$defs");
        Set<List<String>> open = new HashSet<List<String>>();
        for (String[] path : OPEN_CONTENT) {
            open.add(Arrays.asList(path));
        }
        for (String name : MODULE_ARRAYS) {
            open.add(Arrays.asList(name, ELEMENT));
        }
        Map<List<String>, Set<Registration>> found = new LinkedHashMap<List<String>, Set<Registration>>();
        new SchemaWalker(defs, open, found).walk(schema, new ArrayList<String>(), 0);

        Position root = new Position();
        SortedMap<String, Registration> registrations = new TreeMap<String, Registration>();
        for (Map.Entry<List<String>, Set<Registration>> entry : found.entrySet()) {
            String where = displayPath(entry.getKey());
            if (entry.getValue().size() != 1) {
                throw new IllegalStateException("the schema registers " + where + " more than once: " + entry.getValue());
            }
            Registration registration = entry.getValue().iterator().next();
            Position position = root;
            for (String step : entry.getKey()) {
                position = position.child(step);
            }
            position.registration = registration;
            registrations.put(where, registration);
        }
        SortedSet<String> openPaths = new TreeSet<String>();
        for (List<String> path : open) {
            Position position = root;
            for (String step : path) {
                position = position.child(step);
            }
            position.openContent = true;
            openPaths.add(displayPath(path));
        }
        return new ArrayRegistry(root, registrations, openPaths);
    }

    /** {@code $}, then {@code .name} per member and {@code []} per element, as the reference prints paths. */
    static String displayPath(List<String> path) {
        StringBuilder out = new StringBuilder("$");
        for (String step : path) {
            out.append(ELEMENT.equals(step) ? ELEMENT : "." + step);
        }
        return out.toString();
    }

    /** The reference's {@code schema_array_registry}. */
    private static final class SchemaWalker {
        private static final int MAX_DEPTH = 64;
        private static final Set<String> REF_ONLY = new HashSet<String>(Arrays.asList("$ref", "description"));

        private final Map<?, ?> defs;
        private final Set<List<String>> open;
        private final Map<List<String>, Set<Registration>> found;

        SchemaWalker(Map<?, ?> defs, Set<List<String>> open, Map<List<String>, Set<Registration>> found) {
            this.defs = defs;
            this.open = open;
            this.found = found;
        }

        private Object resolve(Object ref) {
            String text = (String) ref;
            if (!text.startsWith("#/$defs/") || !defs.containsKey(text.substring("#/$defs/".length()))) {
                throw new IllegalStateException("unresolvable $ref " + text);
            }
            return defs.get(text.substring("#/$defs/".length()));
        }

        private Object deref(Object node) {
            int guard = 0;
            while (node instanceof Map<?, ?> && ((Map<?, ?>) node).containsKey("$ref")
                    && REF_ONLY.containsAll(((Map<?, ?>) node).keySet())) {
                node = resolve(((Map<?, ?>) node).get("$ref"));
                if (++guard > MAX_DEPTH) {
                    throw new IllegalStateException("$ref cycle");
                }
            }
            return node;
        }

        private void membersOf(Object node, List<Map.Entry<String, Object>> out, int depth) {
            if (depth > MAX_DEPTH) {
                throw new IllegalStateException("schema nesting too deep");
            }
            node = deref(node);
            if (!(node instanceof Map<?, ?>)) {
                return;
            }
            Map<?, ?> map = (Map<?, ?>) node;
            if (map.containsKey("$ref")) {
                membersOf(resolve(map.get("$ref")), out, depth + 1);
            }
            Object properties = map.get("properties");
            if (properties instanceof Map<?, ?>) {
                for (Map.Entry<?, ?> entry : ((Map<?, ?>) properties).entrySet()) {
                    out.add(new java.util.AbstractMap.SimpleImmutableEntry<String, Object>(
                            (String) entry.getKey(), entry.getValue()));
                }
            }
            for (String key : new String[] {"allOf", "anyOf", "oneOf"}) {
                Object subs = map.get(key);
                if (subs instanceof List<?>) {
                    for (Object sub : (List<?>) subs) {
                        membersOf(sub, out, depth + 1);
                    }
                }
            }
            for (String key : new String[] {"then", "else"}) {
                if (map.containsKey(key)) {
                    membersOf(map.get(key), out, depth + 1);
                }
            }
        }

        void walk(Object node, List<String> path, int depth) {
            if (depth > MAX_DEPTH) {
                throw new IllegalStateException("schema nesting too deep at " + displayPath(path));
            }
            node = deref(node);
            if (!(node instanceof Map<?, ?>)) {
                return;
            }
            Map<?, ?> map = (Map<?, ?>) node;
            if ("array".equals(map.get("type")) || map.containsKey("x-sqllr-order")) {
                String where = displayPath(path);
                Object order = map.get("x-sqllr-order");
                if (!(order instanceof String)) {
                    throw new IllegalStateException("the array at " + where + " has no x-sqllr-order");
                }
                Order kind = Order.valueOf((String) order);
                Object key = map.get("x-sqllr-sortKey");
                SortKey sortKey = key == null ? null : SortKey.of(key, where);
                if (kind == Order.SET && sortKey == null) {
                    throw new IllegalStateException("the SET at " + where + " has no x-sqllr-sortKey");
                }
                List<String> at = new ArrayList<String>(path);
                Set<Registration> registrations = found.get(at);
                if (registrations == null) {
                    registrations = new LinkedHashSet<Registration>();
                    found.put(at, registrations);
                }
                registrations.add(new Registration(kind, sortKey));
                if (map.containsKey("items")) {
                    path.add(ELEMENT);
                    walk(map.get("items"), path, depth + 1);
                    path.remove(path.size() - 1);
                }
                return;
            }
            if (open.contains(path)) {
                return;
            }
            List<Map.Entry<String, Object>> members = new ArrayList<Map.Entry<String, Object>>();
            membersOf(map, members, 0);
            for (Map.Entry<String, Object> member : members) {
                if (!"extensions".equals(member.getKey())) {
                    path.add(member.getKey());
                    walk(member.getValue(), path, depth + 1);
                    path.remove(path.size() - 1);
                }
            }
        }
    }
}