/*
 * Copyright 2021 Red Hat, Inc. and/or its affiliates
 * and other contributors as indicated by the @author tags.
 * 
 * Licensed 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.keycloak.models.map.storage.jpa.hibernate.jsonb;

import com.fasterxml.jackson.annotation.JsonAutoDetect;
import com.fasterxml.jackson.annotation.JsonIgnore;
import com.fasterxml.jackson.annotation.JsonInclude;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
import com.fasterxml.jackson.annotation.PropertyAccessor;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.SerializationFeature;
import com.fasterxml.jackson.databind.jsontype.impl.LaissezFaireSubTypeValidator;
import com.fasterxml.jackson.databind.module.SimpleModule;
import com.fasterxml.jackson.databind.node.ObjectNode;
import java.io.IOException;
import java.sql.CallableStatement;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Types;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import org.hibernate.HibernateException;
import org.hibernate.type.AbstractSingleColumnStandardBasicType;
import org.hibernate.type.descriptor.ValueBinder;
import org.hibernate.type.descriptor.ValueExtractor;
import org.hibernate.type.descriptor.WrapperOptions;
import org.hibernate.type.descriptor.java.AbstractTypeDescriptor;
import org.hibernate.type.descriptor.java.JavaTypeDescriptor;
import org.hibernate.type.descriptor.java.MutableMutabilityPlan;
import org.hibernate.type.descriptor.sql.BasicBinder;
import org.hibernate.type.descriptor.sql.BasicExtractor;
import org.hibernate.type.descriptor.sql.SqlTypeDescriptor;
import org.keycloak.models.map.client.MapClientEntity;
import org.keycloak.models.map.client.MapProtocolMapperEntity;
import org.keycloak.models.map.client.MapProtocolMapperEntityImpl;
import org.keycloak.models.map.common.DeepCloner;
import org.keycloak.models.map.common.Serialization.IgnoreUpdatedMixIn;
import org.keycloak.models.map.common.Serialization.IgnoredTypeMixIn;
import org.keycloak.models.map.common.UpdatableEntity;
import org.keycloak.models.map.storage.jpa.client.entity.JpaClientMetadata;
import org.keycloak.models.map.storage.jpa.client.JpaClientMapStorage;

public class JsonbType extends AbstractSingleColumnStandardBasicType<Object> {

    public static final JsonbType INSTANCE = new JsonbType();
    public static final ObjectMapper MAPPER = new ObjectMapper()
            .configure(SerializationFeature.FAIL_ON_EMPTY_BEANS, false)
            .enable(SerializationFeature.INDENT_OUTPUT)
            .setSerializationInclusion(JsonInclude.Include.NON_NULL)
            .setVisibility(PropertyAccessor.ALL, JsonAutoDetect.Visibility.NONE)
            .setVisibility(PropertyAccessor.FIELD, JsonAutoDetect.Visibility.ANY)
            .activateDefaultTyping(new LaissezFaireSubTypeValidator(), ObjectMapper.DefaultTyping.JAVA_LANG_OBJECT, JsonTypeInfo.As.PROPERTY)
            .registerModule(new SimpleModule().addAbstractTypeMapping(MapProtocolMapperEntity.class, MapProtocolMapperEntityImpl.class))
            .addMixIn(UpdatableEntity.class, IgnoreUpdatedMixIn.class)
            .addMixIn(DeepCloner.class, IgnoredTypeMixIn.class)
            .addMixIn(MapClientEntity.class, IgnoredClientFieldsMixIn.class);

    abstract class IgnoredClientFieldsMixIn {
        @JsonIgnore public abstract String getId();
        @JsonIgnore public abstract Map<String, List<String>> getAttributes();
    }

    public JsonbType() {
        super(JsonbSqlTypeDescriptor.INSTANCE, JsonbJavaTypeDescriptor.INSTANCE);
    }

    @Override
    public String getName() {
        return "jsonb";
    }

    private static class JsonbSqlTypeDescriptor implements SqlTypeDescriptor {

        private static final JsonbSqlTypeDescriptor INSTANCE = new JsonbSqlTypeDescriptor();

        @Override
        public int getSqlType() {
            return Types.OTHER;
        }

        @Override
        public boolean canBeRemapped() {
            return true;
        }

        @Override
        public <X> ValueBinder<X> getBinder(JavaTypeDescriptor<X> javaTypeDescriptor) {
            return new BasicBinder<X>(javaTypeDescriptor, this) {
                @Override
                protected void doBind(PreparedStatement st, X value, int index, WrapperOptions options) throws SQLException {
                    st.setObject(index, javaTypeDescriptor.unwrap(value, JsonNode.class, options), getSqlType());
                }

                @Override
                protected void doBind(CallableStatement st, X value, String name, WrapperOptions options) throws SQLException {
                    st.setObject(name, javaTypeDescriptor.unwrap(value, JsonNode.class, options), getSqlType());
                }
            };
        }

        @Override
        public <X> ValueExtractor<X> getExtractor(JavaTypeDescriptor<X> javaTypeDescriptor) {
            return new BasicExtractor<X>(javaTypeDescriptor, this) {
                @Override
                protected X doExtract(ResultSet rs, String name, WrapperOptions options) throws SQLException {
                    return javaTypeDescriptor.wrap(extractJson(rs, name), options);
                }

                @Override
                protected X doExtract(CallableStatement statement, int index, WrapperOptions options) throws SQLException {
                    return javaTypeDescriptor.wrap(extractJson(statement, index), options);
                }

                @Override
                protected X doExtract(CallableStatement statement, String name, WrapperOptions options) throws SQLException {
                    return javaTypeDescriptor.wrap(extractJson(statement, name), options);
                }
            };
        }

        private Object extractJson(ResultSet rs, String name) throws SQLException {
            return rs.getObject(name);
        }

        private Object extractJson(CallableStatement statement, int index) throws SQLException {
            return statement.getObject(index);
        }

        private Object extractJson(CallableStatement statement, String name) throws SQLException {
            return statement.getObject(name);
        }
    }

    private static class JsonbJavaTypeDescriptor extends AbstractTypeDescriptor<Object> {

        private static final JsonbJavaTypeDescriptor INSTANCE = new JsonbJavaTypeDescriptor();

        public JsonbJavaTypeDescriptor() {
            super(Object.class, new MutableMutabilityPlan<Object>() {
                @Override
                protected Object deepCopyNotNull(Object value) {
                    try {
                        return MAPPER.readValue(MAPPER.writerFor(value.getClass()).writeValueAsBytes(value), value.getClass());
                    } catch (IOException e) {
                        throw new HibernateException("unable to deep copy object", e);
                    }
                }
            });
        }

        @Override
        public Object fromString(String json) {
            try {
                ObjectNode tree = MAPPER.readValue(json, ObjectNode.class);
                JsonNode ev = tree.get("entityVersion");
                if (ev == null || ! ev.isInt()) throw new IllegalArgumentException("unable to read entity version from " + json);

                int entityVersion = ev.asInt();

                if (entityVersion > JpaClientMapStorage.SUPPORTED_VERSION + 1) throw new IllegalArgumentException("Incompatible entity version: " + entityVersion + ", supportedVersion: " + JpaClientMapStorage.SUPPORTED_VERSION);

                if (entityVersion < JpaClientMapStorage.SUPPORTED_VERSION) {
                    tree = JpaClientMigration.migrateTreeTo(entityVersion, JpaClientMapStorage.SUPPORTED_VERSION, tree);
                }
                return MAPPER.treeToValue(tree, JpaClientMetadata.class);
            } catch (IOException e) {
                throw new HibernateException("unable to read", e);
            }
        }

        @Override
        public <X> X unwrap(Object value, Class<X> type, WrapperOptions options) {
            if (value == null) return null;

            String stringValue = (value instanceof String) ? (String) value : toString(value);
            try {
                return (X) MAPPER.readTree(stringValue);
            } catch (IOException e) {
                throw new HibernateException("unable to read", e);
            }
        }

        @Override
        public <X> Object wrap(X value, WrapperOptions options) {
            if (value == null) return null;

            return fromString(value.toString());
        }

        @Override
        public String toString(Object value) {
            try {
                return MAPPER.writeValueAsString(value);
            } catch (IOException e) {
                throw new HibernateException("unable to tranform value: " + value + " as String.", e);
            }
        }

        @Override
        public boolean areEqual(Object one, Object another) {
            if (one == another) return true;
            if (one == null || another == null) return Objects.equals(one, another);
            try {
                return MAPPER.readTree(toString(one)).equals(
                       MAPPER.readTree(toString(another)));
            } catch (IOException e) {
                throw new HibernateException("unable to perform areEqual", e);
            }
        }
    }
}
