/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */
package org.apache.tinkerpop.gremlin.structure.io.gryo;

import org.apache.tinkerpop.gremlin.process.traversal.P;
import org.apache.tinkerpop.gremlin.process.traversal.Path;
import org.apache.tinkerpop.gremlin.process.traversal.util.ConnectiveP;
import org.apache.tinkerpop.gremlin.structure.Edge;
import org.apache.tinkerpop.gremlin.structure.Property;
import org.apache.tinkerpop.gremlin.structure.Vertex;
import org.apache.tinkerpop.gremlin.structure.VertexProperty;
import org.apache.tinkerpop.gremlin.structure.util.detached.DetachedEdge;
import org.apache.tinkerpop.gremlin.structure.util.detached.DetachedPath;
import org.apache.tinkerpop.gremlin.structure.util.detached.DetachedProperty;
import org.apache.tinkerpop.gremlin.structure.util.detached.DetachedVertex;
import org.apache.tinkerpop.gremlin.structure.util.detached.DetachedVertexProperty;
import org.apache.tinkerpop.gremlin.structure.util.reference.ReferenceEdge;
import org.apache.tinkerpop.gremlin.structure.util.reference.ReferencePath;
import org.apache.tinkerpop.gremlin.structure.util.reference.ReferenceProperty;
import org.apache.tinkerpop.gremlin.structure.util.reference.ReferenceVertex;
import org.apache.tinkerpop.gremlin.structure.util.reference.ReferenceVertexProperty;
import org.apache.tinkerpop.gremlin.util.function.Lambda;
import org.apache.tinkerpop.shaded.kryo.ClassResolver;
import org.apache.tinkerpop.shaded.kryo.Kryo;
import org.apache.tinkerpop.shaded.kryo.KryoException;
import org.apache.tinkerpop.shaded.kryo.Registration;
import org.apache.tinkerpop.shaded.kryo.io.Input;
import org.apache.tinkerpop.shaded.kryo.io.Output;
import org.apache.tinkerpop.shaded.kryo.util.IdentityObjectIntMap;
import org.apache.tinkerpop.shaded.kryo.util.IntMap;
import org.apache.tinkerpop.shaded.kryo.util.ObjectMap;

import java.net.InetAddress;
import java.nio.ByteBuffer;

import static org.apache.tinkerpop.shaded.kryo.util.Util.getWrapperClass;

/**
 * This mapper implementation of the {@code ClassResolver} helps ensure that all Vertex and Edge concrete classes
 * get properly serialized and deserialized by stripping them of their concrete class name so that they are treated
 * generically. See the {@link #getRegistration(Class)} method for the core of this logic.
 *
 * @author Stephen Mallette (http://stephen.genoprime.com)
 */
public class GryoClassResolver implements ClassResolver {
    static public final byte NAME = -1;

    protected Kryo kryo;

    protected final IntMap<Registration> idToRegistration = new IntMap<>();
    protected final ObjectMap<Class, Registration> classToRegistration = new ObjectMap<>();

    protected IdentityObjectIntMap<Class> classToNameId;
    protected IntMap<Class> nameIdToClass;
    protected ObjectMap<String, Class> nameToClass;
    protected int nextNameId;

    private int memoizedClassId = -1;
    private Registration memoizedClassIdValue;
    private Class memoizedClass;
    private Registration memoizedClassValue;

    @Override
    public void setKryo(Kryo kryo) {
        this.kryo = kryo;
    }

    @Override
    public Registration register(final Registration registration) {
        if (null == registration) throw new IllegalArgumentException("Registration cannot be null.");
        if (registration.getId() != NAME) idToRegistration.put(registration.getId(), registration);

        classToRegistration.put(registration.getType(), registration);
        if (registration.getType().isPrimitive())
            classToRegistration.put(getWrapperClass(registration.getType()), registration);
        return registration;
    }

    @Override
    public Registration registerImplicit(final Class type) {
        return register(new Registration(type, kryo.getDefaultSerializer(type), NAME));
    }

    @Override
    public Registration getRegistration(final Class clazz) {
        // force all instances of Vertex, Edge, VertexProperty, etc. to their respective interface
        final Class type;
        if (!ReferenceVertex.class.isAssignableFrom(clazz) && !DetachedVertex.class.isAssignableFrom(clazz) && Vertex.class.isAssignableFrom(clazz))
            type = Vertex.class;
        else if (!ReferenceEdge.class.isAssignableFrom(clazz) && !DetachedEdge.class.isAssignableFrom(clazz) && Edge.class.isAssignableFrom(clazz))
            type = Edge.class;
        else if (!ReferenceVertexProperty.class.isAssignableFrom(clazz) && !DetachedVertexProperty.class.isAssignableFrom(clazz) && VertexProperty.class.isAssignableFrom(clazz))
            type = VertexProperty.class;
        else if (!ReferenceProperty.class.isAssignableFrom(clazz) && !DetachedProperty.class.isAssignableFrom(clazz) && !DetachedVertexProperty.class.isAssignableFrom(clazz) && Property.class.isAssignableFrom(clazz))
            type = Property.class;
        else if (!ReferencePath.class.isAssignableFrom(clazz) && !DetachedPath.class.isAssignableFrom(clazz) && Path.class.isAssignableFrom(clazz))
            type = Path.class;
        else if (Lambda.class.isAssignableFrom(clazz))
            type = Lambda.class;
        else if (ByteBuffer.class.isAssignableFrom(clazz))
            type = ByteBuffer.class;
        else if (Class.class.isAssignableFrom(clazz))
            type = Class.class;
        else if (InetAddress.class.isAssignableFrom(clazz))
            type = InetAddress.class;
        else if (ConnectiveP.class.isAssignableFrom(clazz))
            type = P.class;
        else
            type = clazz;

        if (type == memoizedClass) return memoizedClassValue;
        final Registration registration = classToRegistration.get(type);
        if (registration != null) {
            memoizedClass = type;
            memoizedClassValue = registration;
        }

        return registration;
    }

    @Override
    public Registration getRegistration(final int classID) {
        return idToRegistration.get(classID);
    }

    @Override
    public Registration writeClass(final Output output, final Class type) {
        if (null == type) {
            output.writeVarInt(Kryo.NULL, true);
            return null;
        }

        final Registration registration = kryo.getRegistration(type);
        if (registration.getId() == NAME)
            writeName(output, type);
        else
            output.writeVarInt(registration.getId() + 2, true);

        return registration;
    }

    protected void writeName(final Output output, final Class type) {
        output.writeVarInt(NAME + 2, true);
        if (classToNameId != null) {
            final int nameId = classToNameId.get(type, -1);
            if (nameId != -1) {
                output.writeVarInt(nameId, true);
                return;
            }
        }
        // Only write the class name the first time encountered in object graph.
        final int nameId = nextNameId++;
        if (classToNameId == null) classToNameId = new IdentityObjectIntMap<>();
        classToNameId.put(type, nameId);
        output.writeVarInt(nameId, true);
        output.writeString(type.getName());
    }

    @Override
    public Registration readClass(final Input input) {
        final int classID = input.readVarInt(true);
        switch (classID) {
            case Kryo.NULL:
                return null;
            case NAME + 2: // Offset for NAME and NULL.
                return readName(input);
        }

        if (classID == memoizedClassId) return memoizedClassIdValue;
        final Registration registration = idToRegistration.get(classID - 2);
        if (registration == null) throw new KryoException("Encountered unregistered class ID: " + (classID - 2));
        memoizedClassId = classID;
        memoizedClassIdValue = registration;
        return registration;
    }

    protected Registration readName(final Input input) {
        final int nameId = input.readVarInt(true);
        if (nameIdToClass == null) nameIdToClass = new IntMap<>();
        Class type = nameIdToClass.get(nameId);
        if (type == null) {
            // Only read the class name the first time encountered in object graph.
            final String className = input.readString();
            type = getTypeByName(className);
            if (type == null) {
                try {
                    type = Class.forName(className, false, kryo.getClassLoader());
                } catch (ClassNotFoundException ex) {
                    throw new KryoException("Unable to find class: " + className, ex);
                }
                if (nameToClass == null) nameToClass = new ObjectMap<>();
                nameToClass.put(className, type);
            }
            nameIdToClass.put(nameId, type);
        }
        return kryo.getRegistration(type);
    }

    protected Class<?> getTypeByName(final String className) {
        return nameToClass != null ? nameToClass.get(className) : null;
    }

    @Override
    public void reset() {
        if (!kryo.isRegistrationRequired()) {
            if (classToNameId != null) classToNameId.clear();
            if (nameIdToClass != null) nameIdToClass.clear();
            nextNameId = 0;
        }
    }
}