package org.bdware.sc.engine.hook;

import com.alibaba.datax.transport.transformer.maskingMethods.cryptology.AESEncryptionImpl;
import com.alibaba.datax.transport.transformer.maskingMethods.cryptology.FormatPreservingEncryptionImpl;
import com.alibaba.datax.transport.transformer.maskingMethods.differentialPrivacy.EpsilonDifferentialPrivacyImpl;
import com.alibaba.datax.transport.transformer.maskingMethods.irreversibleInterference.MD5EncryptionImpl;
import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonPrimitive;
import org.bdware.mockjava.JsonVisitor;

public class MaskVisitor extends JsonVisitor {
    JsonElement root;

    public MaskVisitor(JsonElement ret) {
        root = ret;
    }

    @Override
    public JsonVisitor visitObject(JsonObject mask) {
        if (root.isJsonObject()) {
            JsonObject jo = root.getAsJsonObject();
            for (String key : mask.keySet()) {
                if (jo.has(key)) {
                    //TODO
                    MaskVisitor visitor = new MaskVisitor(jo.get(key));
                    visitor.visit(mask.get(key));
                    jo.add(key, visitor.get());
                }
            }
        }
        return this;
    }

    @Override
    public JsonVisitor visitJsonArray(JsonArray ele) {
        if (root.isJsonArray()) {
            JsonArray array = root.getAsJsonArray();
            for (int i = 0; i < array.size(); i++) {
                MaskVisitor visitor = new MaskVisitor(array.get(i));
                visitor.visit(ele.get(0));
                array.set(i, visitor.get());
            }
        }
        return this;
    }

    public JsonElement get() {
        return root;
    }

    @Override

    public JsonVisitor visitPrimitive(JsonPrimitive primitive) {
        //
        if (primitive.isString()) {
            String method = primitive.getAsString();
            try {
                String result = "";
                //md5不需要参数
                if (method.equals("md5")) {
                    MD5EncryptionImpl masker = new MD5EncryptionImpl();
                    result = masker.execute(root.getAsString());
                } else if (method.equals("aes")) {
                    AESEncryptionImpl masker = new AESEncryptionImpl();
                    result = masker.execute(root.getAsString());
                } else if (method.equals("fpe")) {
                    FormatPreservingEncryptionImpl masker = new FormatPreservingEncryptionImpl();
                    result = masker.execute(root.getAsString());
                }
                //edp需要精度的参数
                else if (method.equals("edp")) {
                    EpsilonDifferentialPrivacyImpl masker = new EpsilonDifferentialPrivacyImpl();
                    double epsilon = 1;

                    result = "" + masker.maskOne(root.getAsDouble(), epsilon);
                } else {
                    result = root.getAsString();
                }
                System.out.println(result);
                root = new JsonPrimitive(result);

            } catch (Exception e) {
                System.out.println(e);
            }
        }


        //String result = masker.execute(primitive.toString());
        //System.out.println(result);

        //root = new JsonPrimitive(root.getAsString().substring(0, 2));
        //https://github.com/guohf/DataX-Masking
        return this;
    }
}
