#!/usr/bin/env python
# vim: ai ts=4 sts=4 et sw=4
#
# Copyright (C) 2016 Samsung
#
#    This program is free software; you can redistribute it and/or
#    modify it under the terms of the GNU General Public License
#    as published by the Free Software Foundation; version 2
#    of the License.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with this program; if not, write to the Free Software
#    Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.
#
"""
This code is called by instance scailing job
"""

import os
import sys
import json
import argparse
from time import sleep
import boto3
import base64
from argparse import ArgumentParser
from datetime import datetime, timedelta

# set default char-set endcoding to utf-8
reload(sys)
sys.setdefaultencoding('utf-8') # pylint: disable-msg=E1101

def unicode_to_str(obj):
    """convert unicode object to str"""

    if isinstance(obj, list):
        return [unicode_to_str(element) for element in obj]
    elif isinstance(obj, dict):
        return {unicode_to_str(key) : unicode_to_str(value) for key, value \
                in obj.iteritems()}
    elif isinstance(obj, unicode):
        return obj.encode('utf-8')
    else:
        return obj

def format_output(data):
    print '\n\n-----BEGIN RC OUTPUT-----'
    sys.stdout.flush()
    print data
    sys.stdout.flush()
    print '-----END RC OUTPUT-----\n\n'
    sys.stdout.flush()
    sleep(1)

def stop_instances(ids):
    assert type(ids) == list
    print 'Stopping %s' % ids
    #TODO: Check state before proceed.
    ec2 = boto3.resource('ec2')
    ec2.instances.filter(InstanceIds=ids).stop()

def start_instances(ids):
    assert type(ids) == list
    print 'Starting %s' % ids
    #TODO: Check state before proceed.
    ec2 = boto3.resource('ec2')
    ec2.instances.filter(InstanceIds=ids).start()

def terminate_instances(ids):
    assert type(ids) == list
    print 'Terminating %s' % ids
    #TODO: Check state before proceed.
    ec2 = boto3.resource('ec2')
    ec2.instances.filter(InstanceIds=ids).terminate()

def execute_command(userdata, ids):
    assert type(ids) == list
    print 'Executing %s' % ids
    #TODO: Check state before proceed.
    ec2 = boto3.client('ssm')
    a = ec2.send_command(InstanceIds=ids,
                         DocumentName='AWS-RunShellScript',
                         Parameters={"commands":["ifconfig"]})
    print '********'
    print a
    print '********'

def describe_instance_by_instanceid(instance_ids):
    if type(instance_ids) == list:
        ec2 = boto3.resource('ec2')
        filtered_ids = ec2.instances.filter(InstanceIds=instance_ids)
    else:
        filtered_ids = instance_ids

    for instance in filtered_ids:
        print '========================'
        print "ami_launch_index = %s" % instance.ami_launch_index
        print "architecture = %s" % instance.architecture
        print "block_device_mappings = %s" % unicode_to_str(instance.block_device_mappings)
        print "client_token = %s" % instance.client_token
        print "ebs_optimized = %s" % instance.ebs_optimized
        print "ena_support = %s" % instance.ena_support
        print "hypervisor = %s" % instance.hypervisor
        print "iam_instance_profile = %s" % instance.iam_instance_profile
        print "image_id = %s" % instance.image_id
        print "instance_id = %s" % instance.instance_id
        print "instance_lifecycle = %s" % instance.instance_lifecycle
        print "instance_type = %s" % instance.instance_type
        print "kernel_id = %s" % instance.kernel_id
        print "key_name = %s" % instance.key_name
        print "launch_time = %s" % instance.launch_time
        print "monitoring = %s" % unicode_to_str(instance.monitoring)
        print "network_interfaces_attribute = %s" % unicode_to_str(instance.network_interfaces_attribute)
        print "placement = %s" % unicode_to_str(instance.placement)
        print "platform = %s" % instance.platform
        print "private_dns_name = %s" % instance.private_dns_name
        print "private_ip_address = %s" % instance.private_ip_address
        print "product_codes = %s" % instance.product_codes
        print "public_dns_name = %s" % instance.public_dns_name
        print "public_ip_address = %s" % instance.public_ip_address
        print "ramdisk_id = %s" % instance.ramdisk_id
        print "root_device_name = %s" % instance.root_device_name
        print "root_device_type = %s" % instance.root_device_type
        print "security_groups = %s" % unicode_to_str(instance.security_groups)
        print "source_dest_check = %s" % instance.source_dest_check
        print "spot_instance_request_id = %s" % instance.spot_instance_request_id
        print "sriov_net_support = %s" % instance.sriov_net_support
        print "state = %s" % unicode_to_str(instance.state)
        print "state_reason = %s" % instance.state_reason
        print "state_transition_reason = %s" % instance.state_transition_reason
        print "subnet_id = %s" % instance.subnet_id
        print "tags = %s" % unicode_to_str(instance.tags)
        print "virtualization_type = %s" % instance.virtualization_type
        print "vpc_id = %s" % instance.vpc_id
        sys.stdout.flush()

def get_essential_info_from_instance(instance_id):
    if type(instance_id) == list:
        ec2 = boto3.resource('ec2')
        filtered_ids = ec2.instances.filter(InstanceIds=instance_id)
    else:
        filtered_ids = instance_id
    ret_data = {}
    for instance in filtered_ids:
        ret_my = {}
        ret_my['instance_id'] = instance.instance_id
        ret_my['private_ip_address'] = instance.private_ip_address
        ret_my['launch_time'] = str(instance.launch_time)
        ret_my['state'] = str(instance.state['Name'])
        ret_my['tags'] = unicode_to_str(instance.tags)
        ret_data[instance.instance_id] = ret_my
    return ret_data

def get_instances_by_tag(tag):
    ret_ids = []
    ec2 = boto3.resource('ec2')
    myfilter = []
    for x in tag:
        myfilter.append({'Name': 'tag:%s' % x['Key'], 'Values': ['%s' % x['Value']]})
    return ec2.instances.filter(Filters=myfilter)

def create_ec2_ondemand_instance(ImageId, MinCount, MaxCount, KeyName, SecurityGroupIds, \
                                 InstanceType, Placement, SubnetId, EbsOptimized, Tags, UserData):
    print 'Creating instance....'
    ec2 = boto3.resource('ec2')

    ret_data = {}

    if False: #TEST-PURPOSE ONLY
        test_inst = ['i-0000', 'i-0000']
        describe_instance_by_instanceid(test_inst)
        sys.stdout.flush()
        ret_data = get_essential_info_from_instance(test_inst)
        format_output(ret_data)
    else:
        print '\nCREATE STARTED AT %s' % (str(datetime.now()))

        rc = ec2.create_instances(
            ImageId=ImageId,
            MinCount=MinCount,
            MaxCount=MaxCount,
            KeyName=KeyName,
            SecurityGroupIds=SecurityGroupIds,
            InstanceType=InstanceType,
            Placement=Placement,
            SubnetId=SubnetId,
            EbsOptimized=EbsOptimized,
            TagSpecifications=[{'ResourceType': 'instance',
                                'Tags': Tags
                               }],
            UserData=UserData
            )
        print '\nCREATE FINISHED AT %s' % (str(datetime.now()))

        generated_ids = [ x.instance_id for x in rc ]
        describe_instance_by_instanceid(generated_ids)
        ret_data = get_essential_info_from_instance(generated_ids)
        format_output(ret_data)
        return rc

    return []

def ping_main(args):
    tags = [ {'Key': t.split(':')[0], 'Value': t.split(':')[1] + "*"} \
             for t in [ x for x in args.tagname.split(',')] ]
    my_ids = get_instances_by_tag(tag=tags)
    describe_instance_by_instanceid(my_ids)
    ret_data = get_essential_info_from_instance(my_ids)
    format_output(ret_data)
    print '\nMAIN FINISHED AT %s' % (str(datetime.now()))

def terminate_main(args):
    instance_ids = args.instanceid.split(',')
    print instance_ids
    #stop_instances(instance_ids)
    terminate_instances(instance_ids)
    describe_instance_by_instanceid(instance_ids)
    print '\nMAIN FINISHED AT %s' % (str(datetime.now()))

def execute_main(args):
    instance_ids = args.instanceid.split(',')
    print instance_ids
    execute_command(args.userdata, instance_ids)
    describe_instance_by_instanceid(instance_ids)
    print '\nMAIN FINISHED AT %s' % (str(datetime.now()))

def create_main(args):
    # All arguments are mandatory
    assert args.amiid
    assert args.mincount
    assert args.maxcount
    assert args.keyname
    assert args.securitygroupids
    assert args.instancetype
    assert args.placement
    assert args.subnetid
    assert args.ebsoptimized
    assert args.tags
    print 'ImageId: [%s]' % args.amiid
    print 'MinCount: [%d]' % int(args.mincount)
    print 'MaxCount: [%d]' % int(args.maxcount)
    print 'KeyName: [%s]' % args.keyname
    print 'SecurityGroupIds: [%s]' % args.securitygroupids.split(',')
    print 'InstanceType: [%s]' % args.instancetype
    placement = dict(x.split(':') for x in args.placement.split(','))
    print 'Placement: [%s]' % placement
    print 'SubnetId: [%s]' % args.subnetid
    print 'EbsOptimized: [%s]' % json.loads(args.ebsoptimized.lower())
    print 'Tags: [%s], type(%s)' % (args.tags, type(args.tags))
    tags = [ {'Key': t.split(':')[0], 'Value': t.split(':')[1]} \
             for t in [ x for x in args.tags.split(',')] ]
    print 'Tags: [%s]' % tags
    userdata = base64.b64decode(args.userdata)
    #print 'UserData: [%s]' % userdata

    create_ec2_ondemand_instance(args.amiid,
                                 int(args.mincount),
                                 int(args.maxcount),
                                 args.keyname,
                                 args.securitygroupids.split(','),
                                 args.instancetype,
                                 placement,
                                 args.subnetid,
                                 json.loads(args.ebsoptimized.lower()),
                                 tags,
                                 userdata)
    print '\nMAIN FINISHED AT %s' % (str(datetime.now()))
    sys.stdout.flush()
    sleep(1)

def argument_parsing(argv):
    """Any arguments passed from user"""

    parser = argparse.ArgumentParser(description='AWS control interface')

    subparsers = parser.add_subparsers(dest='subcommands')

    #### [subcommand - PING] ####
    cmd_ping = subparsers.add_parser('ping')
    cmd_ping.add_argument('-t', '--tagname', action='store', dest='tagname', \
                      help='Comma(,) separeted KEY:VALUE paires of the tag.' \
                           ' ex) Name:myinstancename,hostname:myhost')
    cmd_ping.add_argument('-e', '--env', action='store', dest='env', \
                      help='env. ex) stage')

    #### [subcommand - CREATE] ####
    cmd_create = subparsers.add_parser('create')
    cmd_create.add_argument('-a', '--amiid', action='store', dest='amiid', \
                        help='ImageID. ex) ami-1234abcd')
    cmd_create.add_argument('-n', '--mincount', action='store', default='1', dest='mincount', \
                        help='MinCount(Integer).')
    cmd_create.add_argument('-x', '--maxcount', action='store', default='1', dest='maxcount', \
                        help='MaxCount(Integer).')
    cmd_create.add_argument('-k', '--keyname', action='store', dest='keyname', \
                        help='KeyName.')
    cmd_create.add_argument('-s', '--securitygroupids', action='store', dest='securitygroupids', \
                        help='Comma(,) separeted SecurityGroupIds.' \
                             ' ex) sg-abcd1234,sg-1234abcd')
    cmd_create.add_argument('-i', '--instancetype', action='store', dest='instancetype', \
                        help='InstanceType. ex) t2.micro')
    cmd_create.add_argument('-p', '--placement', action='store', dest='placement', \
                        help='Comma(,) separated KEY:VALUE pairs.' \
                             ' ex) Tenancy:default,GroupName:,AvailabilityZone:ap-northeast-2a')
    cmd_create.add_argument('-b', '--subnetid', action='store', dest='subnetid', \
                        help='SubnetId. ex) subnet-abcd1234')
    cmd_create.add_argument('-e', '--ebsoptimized', action='store', default='false', dest='ebsoptimized', \
                        help='EbsOptimized(True or False)')
    cmd_create.add_argument('-t', '--tags', action='store', dest='tags', \
                        help='Comma(,) separeted KEY:VALUE pairs.' \
                             ' ex) Name:test_name_tag,env:stage,hostname:mycom')
    cmd_create.add_argument('-u', '--userdata', action='store', dest='userdata', \
                        help='UserData(String).')

    #### [subcommand - TERMINATE] ####
    cmd_terminate = subparsers.add_parser('terminate')
    cmd_terminate.add_argument('-i', '--instanceid', action='store', dest='instanceid', \
                        help='Comma(,) separated instance ids.' \
                             ' ex) i-0123456789abcdef0,i-abcdef01234567890')

    #### [subcommand - EXECUTE] ####
    cmd_execute = subparsers.add_parser('execute')
    cmd_execute.add_argument('-i', '--instanceid', action='store', dest='instanceid', \
                        help='Comma(,) separated instance ids.' \
                             ' ex) i-0123456789abcdef0,i-abcdef01234567890')
    cmd_execute.add_argument('-u', '--userdata', action='store', dest='userdata', \
                        help='UserData(String).')

    return parser.parse_args(argv[1:])

def main(argv):
    print '\nMAIN START AT %s' % (str(datetime.now()))
    args = argument_parsing(argv)
    if args.subcommands == 'ping':
        return ping_main(args)
    elif args.subcommands == 'create':
        return create_main(args)
    elif args.subcommands == 'terminate':
        return terminate_main(args)
    elif args.subcommands == 'execute':
        return execute_main(args)
    else:
        print 'Unsopported command %s' % args.subcommands
        return -1

if __name__ == '__main__':
    sys.exit(main(sys.argv))

