/*
 * Copyright (c) 2016. Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 * http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file 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 com.amazonaws.codegen.customization.processors;

import com.amazonaws.codegen.customization.CodegenCustomizationProcessor;
import com.amazonaws.codegen.internal.Utils;
import com.amazonaws.codegen.model.config.customization.ShapeModifier;
import com.amazonaws.codegen.model.config.customization.ShapeModifier_ModifyModel;
import com.amazonaws.codegen.model.intermediate.EnumModel;
import com.amazonaws.codegen.model.intermediate.IntermediateModel;
import com.amazonaws.codegen.model.intermediate.MemberModel;
import com.amazonaws.codegen.model.intermediate.ShapeModel;
import com.amazonaws.codegen.model.service.Member;
import com.amazonaws.codegen.model.service.ServiceModel;
import com.amazonaws.codegen.model.service.Shape;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;

/**
 * This processor handles all the modification on the shape members in the
 * pre-process step and then removes all the excluded shapes in the post-process
 * step.
 *
 * This processor can also modify the names of the enum values generated by the
 * code generator.
 */
final class ShapeModifiersProcessor implements CodegenCustomizationProcessor {

    private final Map<String, ShapeModifier> shapeModifiers;

    private static final String ALL = "*";

    ShapeModifiersProcessor(
            Map<String, ShapeModifier> shapeModifiers) {
        this.shapeModifiers = shapeModifiers;
    }

    @Override
    public void preprocess(ServiceModel serviceModel) {

        if (shapeModifiers == null) return;

        for (Entry<String, ShapeModifier> entry : shapeModifiers.entrySet()) {
            String key = entry.getKey();
            ShapeModifier modifier = entry.getValue();

            if (ALL.equals(key)) {
                for (Shape shape : serviceModel.getShapes().values()) {
                    preprocess_ModifyShapeMembers(serviceModel, shape, modifier);
                }

            } else {
                Shape shape = serviceModel.getShapes().get(key);

                if (shape == null) {
                    throw new IllegalStateException(
                            "ShapeModifier customization found for " + key
                            + ", but this shape doesn't exist in the model!");
                }

                preprocess_ModifyShapeMembers(serviceModel, shape, modifier);
            }
        }
    }

    @Override
    public void postprocess(IntermediateModel intermediateModel) {

        if (shapeModifiers == null) return;
        
        for (Entry<String, ShapeModifier> entry : shapeModifiers.entrySet()) {
            String key = entry.getKey();
            ShapeModifier modifier = entry.getValue();

            if (ALL.equals(key)) continue;

            ShapeModel shapeModel;
            try {
                shapeModel = Utils.findShapeModelByC2jName(
                        intermediateModel, key);
            } catch (IllegalArgumentException e) {
                throw new IllegalStateException(String.format(
                        "Cannot find c2j shape [%s] in the intermediate model when processing " +
                        "customization config shapeModifiers.%s",
                        key, key), e);
            }

            if (modifier.isExcludeShape()) {
                shapeModel.getCustomization().setSkipGeneratingModelClass(true);
                shapeModel.getCustomization().setSkipGeneratingMarshaller(true);
                shapeModel.getCustomization().setSkipGeneratingUnmarshaller(true);
            } else if (modifier.getModify() != null) {
                // Modifies properties of a member in shape or shape enum.
                // This customization currently support modifying enum name
                // and marshall/unmarshall location of a member in the Shape.
                for (Map<String, ShapeModifier_ModifyModel> modifies : modifier.getModify()) {
                    for (Entry<String, ShapeModifier_ModifyModel> memberEntry : modifies.entrySet()) {
                        String enumToModify = memberEntry.getKey();
                        ShapeModifier_ModifyModel modifyModel = memberEntry.getValue();

                        postprocess_ModifyMemberProperty(shapeModel, enumToModify, modifyModel);
                    }
                }
            }
        }
    }

    /**
     * Override name of the enums, marshall/unmarshall location of the
     * members in the given shape model.
     */
    private void postprocess_ModifyMemberProperty(ShapeModel shapeModel, String memberName,
                                                  ShapeModifier_ModifyModel modifyModel) {
        if (modifyModel.getEmitEnumName() != null) {
            EnumModel enumModel = shapeModel.findEnumModelByValue(memberName);
            if (enumModel == null) {
                throw new IllegalStateException(
                        String.format("Cannot find enum [%s] in the intermediate model when processing "
                                + "customization config shapeModifiers.%s", memberName, memberName));
            }
            enumModel.setName(modifyModel.getEmitEnumName());
        }

        if (modifyModel.getMarshallLocationName() != null) {
            MemberModel memberModel =  shapeModel.findMemberModelByC2jName(memberName);
            memberModel.getHttp().setMarshallLocationName(modifyModel.getMarshallLocationName());
        }

        if (modifyModel.getUnmarshallLocationName() != null) {
            MemberModel memberModel =  shapeModel.findMemberModelByC2jName(memberName);
            memberModel.getHttp().setUnmarshallLocationName(modifyModel
                    .getUnmarshallLocationName());
        }

    }

    /**
     * Exclude/modify/inject shape members
     */
    private void preprocess_ModifyShapeMembers(ServiceModel serviceModel, Shape shape, ShapeModifier modifier) {

        if (modifier.getModify() != null) {
            for (Map<String, ShapeModifier_ModifyModel> modifies : modifier.getModify()) {
                for (Entry<String, ShapeModifier_ModifyModel> entry : modifies.entrySet()) {

                    String memberToModify = entry.getKey();
                    ShapeModifier_ModifyModel modifyModel = entry.getValue();

                    doModifyShapeMembers(serviceModel, shape, memberToModify, modifyModel);
                }
            }
        }

        if (modifier.getExclude() != null) {
            for (String memberToExclude : modifier.getExclude()) {

                if (shape.getRequired() != null &&
                    shape.getRequired().contains(memberToExclude)) {
                    throw new IllegalStateException(
                            "ShapeModifier.exclude customization found for "
                                    + memberToExclude
                                    + ", but this member is marked as required in the model!");
                }

                if (shape.getMembers() != null) {
                    shape.getMembers().remove(memberToExclude);
                }
            }
        }

        if (modifier.getInject() != null) {
            for (Map<String, Member> injects : modifier.getInject()) {
                if (shape.getMembers() == null) {
                    shape.setMembers(new HashMap<>());
                }
                shape.getMembers().putAll(injects);
            }
        }

        if (modifier.getInjectEnumValues() != null) {
            if (shape.getEnumValues() == null) {
                throw new IllegalStateException("Attempted to add enum members to a non-enum shape!");
            }
            Set<String> extendValues = new HashSet<>(shape.getEnumValues());
            extendValues.addAll(modifier.getInjectEnumValues());
            shape.setEnumValues(new ArrayList<>(extendValues));
        }
    }

    private void doModifyShapeMembers(ServiceModel serviceModel, Shape shape, String memberToModify,
            ShapeModifier_ModifyModel modifyModel) {

        // Currrenly only supports emitPropertyName which is to rename the member
        if (modifyModel.getEmitPropertyName() != null) {
            Member member = shape.getMembers().remove(memberToModify);

            // if location name is not present, set it to the original name
            // to avoid breaking marshaller code
            if (member.getLocationName() == null) {
                member.setLocationName(memberToModify);
            }

            shape.getMembers().put(modifyModel.getEmitPropertyName(), member);
        }
        if (modifyModel.getEmitAsType() != null) {
            // Must create a shape for the primitive type.
            Shape newShapeForType = new Shape();
            newShapeForType.setType(modifyModel.getEmitAsType());
            final String shapeName = "SDK_" + modifyModel.getEmitAsType();
            serviceModel.getShapes().put(shapeName, newShapeForType);

            shape.getMembers().get(memberToModify).setShape(shapeName);
        }
    }

}
