package com.reyun.saas.common.framework.component;

import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.regions.Regions;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
import com.amazonaws.services.dynamodbv2.document.*;
import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec;
import com.amazonaws.services.dynamodbv2.document.utils.ValueMap;
import com.amazonaws.services.dynamodbv2.model.KeysAndAttributes;
import com.amazonaws.services.dynamodbv2.model.WriteRequest;
import com.google.common.collect.Lists;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.PostConstruct;
import java.util.*;

/**
 * api使用：https://docs.amazonaws.cn/amazondynamodb/latest/developerguide/GettingStarted.Java.01.html
 * CRUD 操作使用Amazon SDK for Java记录 API：https://docs.amazonaws.cn/amazondynamodb/latest/developerguide/JavaDocumentAPICRUDExample.html
 * @date 2016年11月15日
 */
@Component
@Slf4j
public class AwsDynamoDBComponent {
    @Value("${S3.accessKeyId}")
    private String accessKeyId;
    @Value("${S3.accessKeySecret}")
    private String accessKeySecret;
    @Value("${S3.bucketName}")
    private String bucketName;
    @Value("${S3.region}")
    private String regionName;
    
	private  DynamoDB dynamoDB;
	private  AmazonDynamoDBClient client ;

    @PostConstruct
	public void init() {
		if ((accessKeyId != null) && (accessKeySecret != null)) {
			AWSCredentials credentials = new BasicAWSCredentials(accessKeyId, accessKeySecret);
			if (client == null){
				client = new AmazonDynamoDBClient(credentials);
				client.setRegion(com.amazonaws.regions.Region.getRegion(Regions.fromName(regionName)));
			}
			dynamoDB = new DynamoDB(client);
		}
	}

	public DynamoDB getDynamoDB() {
	    return dynamoDB;
	}

	public Table getTable(String table) {
	    return dynamoDB.getTable(table);
	}

    /**
     * 批量写
     */
    public void batchWrite(String tableName, List<Item> items , int batchSize)
    {
        try {
            List<Item> processWriteItems = new ArrayList<Item>();
            Map<String, List<WriteRequest>> unprocessedItems =  new HashMap<>() ;
            for (Item item : items) {
                if (item != null){
                    processWriteItems.add(item);
                }
                if (processWriteItems.size() >= batchSize) {
                    BatchWriteItemOutcome outcome = dynamoDB.batchWriteItem(new TableWriteItems(tableName).withItemsToPut(processWriteItems));
                    // Check for unprocessed keys which could happen if you exceed
                    // provisioned throughput
                    unprocessedItems.putAll(outcome.getUnprocessedItems());
                    processWriteItems.clear();
                }
            }
            if (processWriteItems.size() > 0) {
                BatchWriteItemOutcome outcome2 =dynamoDB.batchWriteItem(new TableWriteItems(tableName).withItemsToPut(processWriteItems));
                // Check for unprocessed keys which could happen if you exceed
                // provisioned throughput
                unprocessedItems.putAll(outcome2.getUnprocessedItems());
                processWriteItems.clear();
            }

            if (!unprocessedItems.isEmpty()) {
                log.warn("BatchWrite  the unprocessed size："+ unprocessedItems.size()+", keys  "+ unprocessedItems.get(tableName));
                BatchWriteItemOutcome outcome3  = dynamoDB.batchWriteItemUnprocessed(unprocessedItems);
                unprocessedItems.putAll(outcome3.getUnprocessedItems());
                if (!unprocessedItems.isEmpty()){
                    log.error("BatchWrite  the unprocessed size："+ unprocessedItems.size()+", keys  "+ unprocessedItems.get(tableName));
                    unprocessedItems.clear();
                }
            }
        }
        catch (Exception e) {
            log.error(e.getMessage(), e);
        }
    }

    /**
     * 批量读
     */
    public List<Item> batchQueryByPrimaryKey(String tableName, String keyName, int batchSize, List<String> values){
        List<Item> result = new ArrayList<>();
        List<List<String>> list = Lists.partition(values, batchSize);
        Map<String, KeysAndAttributes> unprocessed = new HashMap<>();

        for (List<String> vals : list){
            String[] keys = new String[vals.size()];
            keys = vals.toArray(keys);
            TableKeysAndAttributes tableKeysAndAttributes = new TableKeysAndAttributes(tableName).addHashOnlyPrimaryKeys(keyName, keys);
            BatchGetItemOutcome outcome = dynamoDB.batchGetItem(tableKeysAndAttributes);
            result.addAll(outcome.getTableItems().get(tableName));

            // Check for unprocessed keys which could happen if you exceed
            // provisioned
            // throughput or reach the limit on response size.
            unprocessed.putAll(outcome.getUnprocessedKeys());
        }

        if ( !unprocessed.isEmpty()) {
            log.warn("BatchRead  the unprocessed size："+ unprocessed.size()+", keys  "+ unprocessed.get(tableName));
            BatchGetItemOutcome outcome2 = dynamoDB.batchGetItemUnprocessed(unprocessed);
            result.addAll(outcome2.getTableItems().get(tableName));
            unprocessed = outcome2.getUnprocessedKeys();
            if (!unprocessed.isEmpty()){
                log.error("BatchRead  the unprocessed size："+ unprocessed.size()+", keys  "+ unprocessed.get(tableName));
                unprocessed.clear();
            }
        }

        return result;
    }

    /**
     * 向set集合添加数据
     * @param tablename
     * @param keyName
     * @param keyVal
     * @param attributes 如果为嵌套结构，则按顺序添加attribute(attribute如果不存在会出现异常)
     * @param set
     */
    @SuppressWarnings({ "rawtypes", "unchecked" })
    public void addDataToSet(String tablename, String keyName, String keyVal, List<String> attributes, Set set) {
        HashMap<String, String> nameMap = new HashMap<>();
        StringBuffer sb = new StringBuffer();
        for (int i = 0; i < attributes.size(); i++) {
            String attributeName = "#d" + i;
            nameMap.put(attributeName, attributes.get(i));
            if (sb.length() > 0) {
                sb.append(".");
            }
            sb.append(attributeName);
        }
        UpdateItemSpec updateItemSpec = new UpdateItemSpec().withPrimaryKey(keyName, keyVal)
                .withUpdateExpression("ADD " + sb.toString() + " :d ").withNameMap(nameMap)
                .withValueMap(new ValueMap().withStringSet(":d", set));
        dynamoDB.getTable(tablename).updateItem(updateItemSpec);
    }

    /**
     * 从set集合删除数据
     * @param tablename
     * @param keyName
     * @param keyVal
     * @param attributes 如果为嵌套结构，则按顺序添加attribute(attribute如果不存在会出现异常)
     * @param set
     */
    @SuppressWarnings({ "rawtypes", "unchecked" })
    public void deleteDataFromSet(String tablename, String keyName, String keyVal, List<String> attributes, Set set) {
        HashMap<String, String> nameMap = new HashMap<>();
        StringBuffer sb = new StringBuffer();
        for (int i = 0; i < attributes.size(); i++) {
            String attributeName = "#d" + i;
            nameMap.put(attributeName, attributes.get(i));
            if (sb.length() > 0) {
                sb.append(".");
            }
            sb.append(attributeName);
        }
        UpdateItemSpec updateItemSpec = new UpdateItemSpec().withPrimaryKey(keyName, keyVal)
                .withUpdateExpression("DELETE  " + sb.toString() + " :d ").withNameMap(nameMap)
                .withValueMap(new ValueMap().withStringSet(":d", set));
        dynamoDB.getTable(tablename).updateItem(updateItemSpec);
    }

    public PutItemOutcome writeItem(String tableName , Item item){
        return dynamoDB.getTable(tableName).putItem(item);
    }

    public DeleteItemOutcome deleteItem(String tableName, String primaryKeyName, String primaryKeyValue) {
        return dynamoDB.getTable(tableName).deleteItem(primaryKeyName, primaryKeyValue);
    }

    public UpdateItemOutcome updateItem(String tableName , UpdateItemSpec updateItemSpec){
        return dynamoDB.getTable(tableName).updateItem(updateItemSpec);
    }

    public Item queryByPrimaryKey(String tableName, String primaryKeyName, String primaryKeyValue) {
        return dynamoDB.getTable(tableName).getItem(primaryKeyName, primaryKeyValue);
    }

    public void increase(String tableName, String primaryKeyName, String primaryKeyValue, String field, int incValue){
        UpdateItemSpec updateItemSpec = new UpdateItemSpec()
                .withPrimaryKey(primaryKeyName, primaryKeyValue)
                .withUpdateExpression("set " + field + " = " + field + " + :val")
                .withValueMap(new ValueMap().withNumber(":val", incValue));

        dynamoDB.getTable(tableName).updateItem(updateItemSpec);
    }

    public void writeStringSet(String tableName, String primaryKeyName, String primaryKeyValue, String field, Set<String> set) {
        dynamoDB.getTable(tableName).putItem(new Item().withPrimaryKey(primaryKeyName, primaryKeyValue).withStringSet(field, set));
    }

    /**
     *  获取ddb中set集合数据
     * @param tablename  ddb表名
     * @param keyName    key名称
     * @param keyVal    key值
     * @param setName    set名称
     */
    public Set<String> getSetFromDDB(String tablename,String keyName,String keyVal , String setName){
        Set<String> set = new HashSet<>();
        Item sizeItem = this.getTable(tablename).getItem(keyName, keyVal);
        if (sizeItem!=null) {
            set.addAll(sizeItem.getStringSet(setName));
        }
        return set;
    }

    @SuppressWarnings("rawtypes")
    public void writeList(String tableName, String primaryKeyName, String primaryKeyValue, String field, List list) {
        dynamoDB.getTable(tableName).putItem(new Item().withPrimaryKey(primaryKeyName, primaryKeyValue).withList(field, list));
    }

    @SuppressWarnings({ "unchecked", "rawtypes" })
    public List getList(String tablename, String keyName, String keyVal, String listFieldName) {
        List list = new ArrayList<>();
        Item sizeItem = this.getTable(tablename).getItem(keyName, keyVal);
        if (sizeItem!=null) {
            list.addAll(sizeItem.getList(listFieldName));
        }
        return list;
    }

    @SuppressWarnings({ "rawtypes", "unchecked" })
    public void writeMap(String tablename, String primaryKeyName, String primaryKeyValue, String field, Map map) {
        this.getTable(tablename)
                .putItem(new Item().withPrimaryKey(primaryKeyName, primaryKeyValue).withMap(field, map));
    }

    @SuppressWarnings({ "unchecked", "rawtypes" })
    public Map getMap(String tablename, String keyName, String keyVal, String mapFieldName) {
        Map map = new HashMap<>();
        Item item = this.getTable(tablename).getItem(keyName, keyVal);
        if (item!=null) {
            map.putAll(item.getMap(mapFieldName));
        }
        return map;
    }
}
