package com.tzld.piaoquan.ad.engine.commons.score.model;


import com.aliyun.oss.OSS;
import com.aliyun.oss.OSSClientBuilder;
import com.aliyun.oss.common.auth.CredentialsProvider;
import com.aliyun.oss.common.auth.DefaultCredentialProvider;
import com.aliyun.oss.model.OSSObject;
import com.ctrip.framework.apollo.Config;
import com.ctrip.framework.apollo.ConfigService;
import lombok.extern.slf4j.Slf4j;

import java.io.IOException;
import java.io.InputStreamReader;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;


@Slf4j
public class ModelManager {
    private static final int SCHEDULE_PERIOD = 10;
    private static final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1);
    private static ModelManager instance;
    Map<String, ModelLoadTask> loadTasks = new HashMap<>();
    Map<String, String> modelPathMap = new HashMap<>();
    private OSS client;
    private String bucketName;

    private final String modelOssEndpoint = "model.oss.internal.endpoint";
    private final String modelOssAccessKeyId = "model.oss.accessKeyId";
    private final String modelOssAccessKeySecret = "model.oss.accessKetSecret";
    private final String modelOssBucketName = "model.oss.bucketName";

    private ModelManager() {
        // config load
        Config config = ConfigService.getAppConfig();
        String endpoint = config.getProperty(modelOssEndpoint, "");
        String accessKeyId = config.getProperty(modelOssAccessKeyId, "");
        String accessKetSecret = config.getProperty(modelOssAccessKeySecret, "");
        // oss client
        CredentialsProvider credentialsProvider = new DefaultCredentialProvider(accessKeyId, accessKetSecret);
        this.client = new OSSClientBuilder().build(endpoint, credentialsProvider);
        this.bucketName = config.getProperty(modelOssBucketName, "");

        config.addChangeListener(changeEvent -> {
            if (changeEvent.isChanged(modelOssEndpoint)
                    || changeEvent.isChanged(modelOssAccessKeyId)
                    || changeEvent.isChanged(modelOssAccessKeySecret)) {
                String endpointNew = config.getProperty(modelOssEndpoint, "");
                String accessKeyIdNew = config.getProperty(modelOssAccessKeyId, "");
                String accessKetSecretNew = config.getProperty(modelOssAccessKeySecret, "");
                CredentialsProvider credentialsProviderNew = new DefaultCredentialProvider(accessKeyIdNew,
                        accessKetSecretNew);
                this.client = new OSSClientBuilder().build(endpointNew, credentialsProviderNew);
            }
            if (changeEvent.isChanged(modelOssBucketName)) {
                this.bucketName = config.getProperty(modelOssBucketName, "");
            }
        });


        start(SCHEDULE_PERIOD);
    }

    public static ModelManager getInstance() {
        if (instance == null) {
            synchronized (ModelManager.class) {
                if (instance == null) {
                    instance = new ModelManager();
                }
            }
        }
        return instance;
    }

    /**
     * 添加一个加载任务到管理器
     *
     * @param modelName  Model的名字, 注册到ModelManager的不同model需要不同的名字
     * @param path       Model在OSS上的全路径
     * @param modelClass Model的子类型
     */
    public void registerModel(String modelName, String path, Class<? extends Model> modelClass) throws ModelRegisterException, IOException {
        if (modelPathMap.containsKey(modelName)) {
            return;
            // fail fast
            // throw new RuntimeException(modelName + " already exists");
//            String oldPath = modelPathMap.get(modelName);
//            if (path.equals(oldPath)) {
//                //如果模型的path没有发生改变, 不做任何操作
//                log.info("Model [{}] and Path [{}] has exist", modelName, path);
//                return;
//            } else {
//                //如果模型的path发生改变, 需要注销掉原有的任务
//                unRegisterModel(modelName);
//            }
        }

        modelPathMap.put(modelName, path);
        if (loadTasks.containsKey(path)) {
            ModelLoadTask loadTask = loadTasks.get(path);
            loadTask.refCount++;
        } else {
            ModelLoadTask task = new ModelLoadTask(path, modelClass);
            task.refCount++;
            loadTasks.put(path, task);
            loadModel(task, false, true);
        }
    }

    /**
     * 删除一个加载任务
     *
     * @param modelName Model的名字, 需要和registerModel的名字一致
     */
    private void unRegisterModel(String modelName) {
        if (modelPathMap.containsKey(modelName)) {
            String path = modelPathMap.get(modelName);
            if (loadTasks.containsKey(path)) {
                ModelLoadTask task = loadTasks.get(path);
                task.refCount--;
                if (task.refCount == 0) {
                    loadTasks.remove(path);
                }
            }
            modelPathMap.remove(modelName);
        }
    }

    /**
     * @param modelName
     * @return
     */
    public Model getModel(String modelName) {
        if (modelPathMap.containsKey(modelName) && loadTasks.containsKey(modelPathMap.get(modelName))) {
            return loadTasks.get(modelPathMap.get(modelName)).model;
        } else {
            return null;
        }
    }

    /**
     * 开始调度
     *
     * @param period
     */
    protected void start(long period) {
        final Runnable task = new Runnable() {
            public void run() {
                // 模型更新开关
                // boolean modelUpdateSwitch = Configuration.getBoolean("recommend-service-framework.model-update-switch", true);
                boolean modelUpdateSwitch = true;
                log.info("model update switch [{}]", modelUpdateSwitch);
                if (modelUpdateSwitch) {
                    updateModels(false);
                }
            }
        };
        scheduler.scheduleAtFixedRate(task, 10, period, TimeUnit.MINUTES); // 10分钟
    }

    /**
     * 更新模型
     */
    public void updateModels(final boolean isForceLoads) {
        log.info("begin to update: [{}]", loadTasks.keySet().size());
        for (String modelPath : loadTasks.keySet()) {
            log.debug("load task model path [{}]", modelPath);
            ModelLoadTask task = loadTasks.get(modelPath);
            loadModel(task, isForceLoads, false);
        }
    }

    /**
     * 检查并加载模型
     * <p>
     * 从oss加载：
     * https://help.aliyun.com/zh/oss/developer-reference/streaming-download-7?spm=a2c4g.11186623.0.0.4b527c7dm8LejC
     *
     * @param loadTask
     */
    private void loadModel(final ModelLoadTask loadTask, final boolean isForceLoads, final boolean isRegister) {
        if (loadTask.isLoading) {
            return;
        }
        loadTask.isLoading = true;
        OSSObject ossObj = null;
        try {
            ossObj = client.getObject(bucketName, loadTask.path);
            long timeStamp = ossObj.getObjectMetadata().getLastModified().getTime();
            if (loadTask.lastModifyTime <= timeStamp || isForceLoads) {
                log.info("model file [{}] changed, ready to update, last modify: [{}], local timestamp: {}", loadTask.path, loadTask.lastModifyTime, timeStamp);

                Model model = loadTask.modelClass.newInstance();
                if (model.loadFromStream(ossObj.getObjectContent())) {
                    loadTask.model = model;
                    loadTask.lastModifyTime = timeStamp;
                }

                log.info("model file [{}] update success, last modify: [{}], local timestamp: {}", loadTask.path, loadTask.lastModifyTime, timeStamp);

            }
            ossObj.close();
        } catch (Exception e) {
            log.error("update model fail: {}", loadTask.path, e);
        } finally {
            loadTask.isLoading = false;
            if (ossObj != null) {
                try {
                    ossObj.close();
                } catch (IOException e) {
                    log.error("close ossObj fail", e);
                }
            }
        }
    }

    public class ModelRegisterException extends Exception {

        public ModelRegisterException(String s) {
            super(s);
        }
    }

    /**
     * 调度的任务单元
     */
    private class ModelLoadTask {

        private int refCount;
        private String path;
        private long lastModifyTime;
        private boolean isLoading;
        private Class<? extends Model> modelClass;
        private Model model;

        ModelLoadTask(String path, Class<? extends Model> modelClass) {
            this.refCount = 0;
            this.path = path;
            this.lastModifyTime = 0;
            this.modelClass = modelClass;
        }
    }
}
