/*
* Copyright 2014-2015 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.codesamples.document;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;

import com.amazonaws.AmazonServiceException;
import com.amazonaws.auth.profile.ProfileCredentialsProvider;
import com.amazonaws.services.dynamodbv2.AmazonDynamoDBClient;
import com.amazonaws.services.dynamodbv2.document.DynamoDB;
import com.amazonaws.services.dynamodbv2.document.Item;
import com.amazonaws.services.dynamodbv2.document.ItemCollection;
import com.amazonaws.services.dynamodbv2.document.ScanOutcome;
import com.amazonaws.services.dynamodbv2.document.Table;
import com.amazonaws.services.dynamodbv2.document.spec.ScanSpec;
import com.amazonaws.services.dynamodbv2.model.AttributeDefinition;
import com.amazonaws.services.dynamodbv2.model.KeySchemaElement;
import com.amazonaws.services.dynamodbv2.model.KeyType;
import com.amazonaws.services.dynamodbv2.model.ProvisionedThroughput;

public class DocumentAPIParallelScan {

    // total number of sample items 
    static int scanItemCount = 300;
    
    // number of items each scan request should return
    static int scanItemLimit = 10;
    
    // number of logical segments for parallel scan
    static int parallelScanThreads = 16;
    
    // table that will be used for scanning
    static String parallelScanTestTableName = "ParallelScanTest";
    
    static DynamoDB dynamoDB = new DynamoDB(
        new AmazonDynamoDBClient(new ProfileCredentialsProvider()));

    public static void main(String[] args) throws Exception {
        try {
 
            // Clean up the table
           deleteTable(parallelScanTestTableName);
           createTable(parallelScanTestTableName, 10L, 5L, "Id", "N");
            
            // Upload sample data for scan
           uploadSampleProducts(parallelScanTestTableName, scanItemCount);
            
            // Scan the table using multiple threads
            parallelScan(parallelScanTestTableName, scanItemLimit, parallelScanThreads);
        }  
        catch (AmazonServiceException ase) {
            System.err.println(ase.getMessage());
        }  
    }

    
   
    private static void parallelScan(String tableName, int itemLimit, int numberOfThreads) {
        System.out.println("Scanning " + tableName + " using " + numberOfThreads 
            + " threads " + itemLimit + " items at a time");
        ExecutorService executor = Executors.newFixedThreadPool(numberOfThreads);
        
        // Divide DynamoDB table into logical segments
        // Create one task for scanning each segment
        // Each thread will be scanning one segment
        int totalSegments = numberOfThreads;
        for (int segment = 0; segment < totalSegments; segment++) {
            // Runnable task that will only scan one segment
            ScanSegmentTask task = new ScanSegmentTask(tableName, itemLimit, totalSegments, segment);
            
            // Execute the task
            executor.execute(task);
        }

        shutDownExecutorService(executor); 
    }

    // Runnable task for scanning a single segment of a DynamoDB table
    private static class ScanSegmentTask implements Runnable {
        
        // DynamoDB table to scan
        private String tableName;
        
        // number of items each scan request should return
        private int itemLimit;
        
        // Total number of segments
        // Equals to total number of threads scanning the table in parallel
        private int totalSegments;
        
        // Segment that will be scanned with by this task
        private int segment;
        
        public ScanSegmentTask(String tableName, int itemLimit, int totalSegments, int segment) {
            this.tableName = tableName;
            this.itemLimit = itemLimit;
            this.totalSegments = totalSegments;
            this.segment = segment;
        }
        
        @Override
        public void run() {
            System.out.println("Scanning " + tableName + " segment " + segment + " out of " + totalSegments + " segments " + itemLimit + " items at a time...");
            int totalScannedItemCount = 0;

            Table table = dynamoDB.getTable(tableName);
            
            try {
                ScanSpec spec = new ScanSpec()
                    .withMaxResultSize(itemLimit)
                    .withTotalSegments(totalSegments)
                    .withSegment(segment);
                
                ItemCollection<ScanOutcome> items = table.scan(spec);
                Iterator<Item> iterator = items.iterator();
                  
                Item currentItem = null;
                while (iterator.hasNext()) {
                    totalScannedItemCount++;
                    currentItem = iterator.next();
                    System.out.println(currentItem.toString());
                }    
                    
            } catch (Exception e) {
                System.err.println(e.getMessage());
            } finally {
                System.out.println("Scanned " + totalScannedItemCount 
                    + " items from segment " + segment + " out of " 
                    + totalSegments + " of " + tableName);
            }
        }
    }
    

    private static void uploadSampleProducts(String tableName, int itemCount) {
        System.out.println("Adding " + itemCount +  " sample items to " + tableName);
        for (int productIndex = 0; productIndex < itemCount; productIndex++) {
            uploadProduct(tableName, productIndex); 
        }
    }

    private static void uploadProduct(String tableName, int productIndex) {

        Table table = dynamoDB.getTable(tableName);

        try {
            System.out.println("Processing record #" + productIndex);

            Item item = new Item()
                .withPrimaryKey("Id", productIndex)
                .withString("Title", "Book " + productIndex + " Title")
                .withString("ISBN", "111-1111111111")
                .withStringSet(
                    "Authors",
                    new HashSet<String>(Arrays.asList("Author1")))
                .withNumber("Price", 2)
                .withString("Dimensions", "8.5 x 11.0 x 0.5")
                .withNumber("PageCount", 500)
                .withBoolean("InPublication", true)
                .withString("ProductCategory", "Book");
            table.putItem(item);

        }   catch (Exception e) {
            System.err.println("Failed to create item " + productIndex + " in " + tableName);
            System.err.println(e.getMessage());
        }
    }

    private static void deleteTable(String tableName){
        try {
            
            Table table = dynamoDB.getTable(tableName);
            table.delete();
            System.out.println("Waiting for " + tableName
                + " to be deleted...this may take a while...");
            table.waitForDelete();
               
        } catch (Exception e) {
            System.err.println("Failed to delete table " + tableName);
            e.printStackTrace(System.err);
        }
    }
    
    private static void createTable(
        String tableName, long readCapacityUnits, long writeCapacityUnits,
        String hashKeyName, String hashKeyType) {
        
        createTable(tableName, readCapacityUnits, writeCapacityUnits, 
           hashKeyName,  hashKeyType, null, null);    
    }
    
    private static void createTable(
        String tableName, long readCapacityUnits, long writeCapacityUnits,
        String hashKeyName, String hashKeyType, 
        String rangeKeyName, String rangeKeyType) {
        
        try {
            System.out.println("Creating table " + tableName);
            
            List<KeySchemaElement> keySchema = new ArrayList<KeySchemaElement>();
            keySchema.add(new KeySchemaElement()
                .withAttributeName(hashKeyName)
                .withKeyType(KeyType.HASH));
            
            List<AttributeDefinition> attributeDefinitions = new ArrayList<AttributeDefinition>();
            attributeDefinitions.add(new AttributeDefinition()
                .withAttributeName(hashKeyName)
                .withAttributeType(hashKeyType));

            if (rangeKeyName != null){
                keySchema.add(new KeySchemaElement()
                    .withAttributeName(rangeKeyName)
                    .withKeyType(KeyType.RANGE));
                attributeDefinitions.add(new AttributeDefinition()
                      .withAttributeName(rangeKeyName)
                      .withAttributeType(rangeKeyType));
            }

            Table table = dynamoDB.createTable(tableName, 
                keySchema,
                attributeDefinitions, 
                new ProvisionedThroughput()
                    .withReadCapacityUnits(readCapacityUnits)
                    .withWriteCapacityUnits(writeCapacityUnits));
            System.out.println("Waiting for " + tableName
                + " to be created...this may take a while...");
            table.waitForActive();
       
            
        } catch (Exception e) {
            System.err.println("Failed to create table " + tableName);
            e.printStackTrace(System.err);
        }
    }

    private static void shutDownExecutorService(ExecutorService executor) {
        executor.shutdown();
        try {
            if (!executor.awaitTermination(10, TimeUnit.SECONDS)) {
                executor.shutdownNow();
            }
        } catch (InterruptedException e) {
            executor.shutdownNow();
            
            // Preserve interrupt status
            Thread.currentThread().interrupt();
        }
    }
}