# Copyright 2015 Adaptavist.com Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License 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.

require_relative './cloud_connection.rb'
using Rainbow
module AvstCloud
    
    class AwsConnection < AvstCloud::CloudConnection
        
        attr_accessor :region

        def initialize(provider_user, provider_pass, region)
            super('aws', provider_user, provider_pass)
            @region = region
        end
        
        def server(server_name, root_user, root_password, os=nil)
            server = find_fog_server(server_name)
            if !root_user and os
                root_user = user_from_os(os)
            end
            AvstCloud::AwsServer.new(server, server_name, server.public_ip_address, root_user, root_password)
        end

        def create_server(server_name, flavour, os, key_name, ssh_key, subnet_id, security_group_ids, ebs_size, hdd_device_path, ami_image_id, availability_zone, additional_hdds={}, vpc=nil, created_by=nil, custom_tags={}, root_username=nil, create_elastic_ip=false, encrypt_root=false ,root_encryption_key=nil, delete_root_disk=true, root_disk_type='gp2', root_disk_iops=0)
            # Permit named instances from DEFAULT_FLAVOURS
            flavour = flavour || "t2.micro"
            os = os || "ubuntu-14"
            ami_image_id = ami_image_id || "ami-f0b11187"
            device_name = hdd_device_path || '/dev/sda1'

            root_user = root_username || user_from_os(os)
            unless File.file?(ssh_key)
                logger.error "Could not find local SSH key '#{ssh_key}'".red
                raise "Could not find local SSH key '#{ssh_key}'"
            end

            existing_servers    = all_named_servers(server_name)
            restartable_servers = existing_servers.select{ |serv| serv.state == 'stopped' }
            running_servers     = existing_servers.select{ |serv| serv.state != 'stopped' && serv.state != 'terminated' }

            if running_servers.length > 0
                running_servers.each do |server|
                    logger.error "Server #{server_name} with id #{server.id} found in state: #{server.state}".yellow
                end
                raise "Server with the same name found!"

            elsif restartable_servers.length > 0
                if restartable_servers.length > 1
                    running_servers.each do |server|
                        logger.error "Server #{server_name} with id #{server.id} found in state: #{server.state}. Can not restart".yellow
                    end
                    raise "Too many servers can be restarted."
                end
                server = restartable_servers.first
                server.start
                result_server = AvstCloud::AwsServer.new(server, server_name, server.public_ip_address, root_user, ssh_key)
                result_server.wait_for_state() {|serv| serv.ready?}
                logger.debug "[DONE]\n\n"
                logger.debug "The server was successfully re-started.\n\n"
                result_server
            else
                logger.debug "Creating EC2 server:"
                logger.debug "Server name        - #{server_name}"
                logger.debug "Operating system   - #{os}"
                logger.debug "image_template_id  - #{ami_image_id}"
                logger.debug "flavour            - #{flavour}"
                logger.debug "key_name           - #{key_name}"
                logger.debug "ssh_key            - #{ssh_key}"
                logger.debug "root user          - #{root_user}"
                logger.debug "subnet_id          - #{subnet_id}"
                logger.debug "security_group_ids - #{security_group_ids}"
                logger.debug "region             - #{@region}"
                logger.debug "availability_zone  - #{availability_zone}"
                logger.debug "ebs_size           - #{ebs_size}"
                logger.debug "hdd_device_path    - #{device_name}"
                logger.debug "additional_hdds    - #{additional_hdds}"
                logger.debug "vpc                - #{vpc}"

                create_ebs_volume = nil
                if ebs_size
                    # in case of centos ami we need to use /dev/xvda this is ami dependent
                    root_disk = { 
                        :DeviceName => device_name,
                        'Ebs.VolumeType' => root_disk_type,
                        'Ebs.VolumeSize' => ebs_size,
                    } 
                    # if the root disk is to be encrypted set te "Encrypted" flag to true, if there is an optional KMS Key ID send that,
                    # if not set to nil and thereby defalt to the default key for EBS
                    if encrypt_root
                        root_disk.merge!('Ebs.Encrypted' => true, 'Ebs.KmsKeyId' => root_encryption_key||nil )
                    end

                    # if we do not want to delete the root disk with the VM set the flag
                    if delete_root_disk == false || delete_root_disk == 'false'
                        root_disk.merge!('Ebs.DeleteOnTermination' => false)
                    end

                    # if this is a provisioned IOPS disk set the iops value
                    if root_disk_type == 'io1'
                        root_disk.merge!('Ebs.Iops' => root_disk_iops)
                    end
                    # add the root disk as the first entry in the array of disks to create/attach
                    create_ebs_volume = [ root_disk ] 

                    if additional_hdds and additional_hdds.is_a?(Hash)
                        additional_hdds.each_value do |disk|
                            volume_type = disk['volume_type'] || 'gp2'
                            if disk['device_name'] && disk['ebs_size']
                                disk_hash = {
                                    :DeviceName => disk['device_name'],
                                    'Ebs.VolumeType' => volume_type,
                                    'Ebs.VolumeSize' => disk['ebs_size']
                                }
                                # if the additional disk is to be encrypted set te "Encrypted" flag to true, if there is an optional KMS Key ID send that,
                                # if not set to nil and thereby defalt to the default key for EBS
                                if disk['encrypted']
                                    disk_hash.merge!('Ebs.Encrypted' => true, 'Ebs.KmsKeyId' => disk['encryption_key_id'] || nil)
                                end

                                # if we do not want to delete the additional disk with the VM set the flag
                                if disk['delete_disk_with_vm'] == false || disk['delete_disk_with_vm'] == 'false'
                                    disk_hash.merge!('Ebs.DeleteOnTermination' => false)
                                end

                                # if the additional disk is an provisioned IOPS disk set the iops value
                                if volume_type == 'io1'
                                    disk_hash.merge!('Ebs.Iops' => disk['volume_iops'] || 0)
                                end

                                # add this disk to the array of all disks to create/attach
                                create_ebs_volume << disk_hash
                            else
                                logger.warn "Failed to create additional hdd, required params device_name (e.g. /dev/sda1) or ebs_size missing: #{disk}"
                            end 
                        end 
                    end
                end

                tags = {
                    'Name' => server_name,
                    'os' => os
                }
                if created_by 
                    tags['created_by'] = created_by
                end
                tags.merge!(custom_tags)

                # create server
                server = connect.servers.create :tags => tags,
                                                :flavor_id => flavour,
                                                :image_id => ami_image_id,
                                                :key_name => key_name,
                                                :subnet_id => subnet_id,
                                                :associate_public_ip => true,
                                                :security_group_ids => security_group_ids,
                                                :availability_zone => availability_zone,
                                                :block_device_mapping => create_ebs_volume,
                                                :vpc => vpc
                
                result_server = AvstCloud::AwsServer.new(server, server_name, nil, root_user, ssh_key)
                # result_server.logger = logger
                # Check every 5 seconds to see if server is in the active state (ready?).
                # If the server has not been built in 5 minutes (600 seconds) an exception will be raised.
                result_server.wait_for_state() {|serv| serv.ready?}

                logger.debug "[DONE]\n\n"

                # create Elastic IP Address if required
                if create_elastic_ip
                    logger.debug("Attempting to create elastic IP address")
                    elastic_ip = connect.allocate_address("vpc").body
                    elastic_ip_address = elastic_ip['publicIp']
                    # if we have a server id and an Elastic public IP attempt to join the two togehter
                    if server.id and elastic_ip_address
                        logger.debug ("Elastic IP #{elastic_ip_address} created, attempting to allocate to server")
                        connect.associate_address(server.id, elastic_ip_address)
                        # reacquire server object as IP has, probably, changed
                        server = find_fog_server(server_name)

                        # create tag on the Elastic IP 
                        # TODO: add ability for other tags to be defined by the user
                        logger.debug("Creating tags on Elastic IP Address #{elastic_ip}\n\n")
                        connect.tags.create(:resource_id => elastic_ip['allocationId'], :key => "Name", :value => server_name)
                    else
                        logger.warn("Elastic IP creation failed, proceeding with non Elastic IP\n\n")
                    end
                end

                logger.debug "The server has been successfully created, to login onto the server:\n"
                logger.debug "\t ssh -i #{ssh_key} #{root_user}@#{server.public_ip_address}\n"
                if create_ebs_volume
                    logger.debug("Creating tags on ebs volumes")
                    ebs_volumes = server.block_device_mapping
                    logger.debug("Creating tags on ebs volumes #{ebs_volumes}")
                    ebs_volumes.each do |ebs_volume|
                        if ebs_volume['volumeId']
                            tags.each do |key, value|
                                connect.tags.create(:resource_id => ebs_volume['volumeId'], :key => key, :value => value)
                            end
                        end
                    end
                end
                result_server.ip_address = server.public_ip_address
                result_server
            end
        end

        def server_status(server_name)
            servers = all_named_servers(server_name).select{|serv| serv.state != 'terminated'}
            if servers.length > 0
                servers.each do |server|
                    logger.debug "Server #{server.id} with name '#{server_name}' exists and has state: #{server.state}"
                end
            else
                logger.debug "Server not found for name: #{server_name}"
            end
        end

        def list_flavours
            connect.flavors.each do |fl|
                logger.debug fl.inspect
            end
        end

        def list_images
            connect.images.each do |im|
                logger.debug im.inspect
            end
        end
        
        # Returns list of servers from fog
        def list_known_servers
            connect.servers.all
        end

        def find_fog_server(server_name, should_fail=true)
            servers = all_named_servers(server_name).select{|serv| serv.state != 'terminated'}
            unless servers.length == 1    
                logger.debug "Found #{servers.length} servers for name: #{server_name}".yellow
                if should_fail
                    raise "Found #{servers.length} servers for name: #{server_name}"
                end
            end
            servers.first
        end

        def delete_elastic_ip(ip_address)
            address = is_elastic_ip(ip_address)
            if address
                logger.debug "Found Elastic IP #{address.public_ip}, attempting to delete"
                logger.debug "Elastic IP #{ip_address} deleted" if address.destroy 
                return true
            else
                logger.debug "IP #{ip_address} does NOT appear to be an Elastic IP"
            end
            return false
        end

    private
        def user_from_os(os)
            case os.to_s
                when /^ubuntu/
                    user = "ubuntu"
                when /^debian/
                    user = "admin"
                when /^centos/
                    user = "ec2-user"
                when /^redhat/
                    user = "ec2-user"
                else
                    user = "root"
            end
            user
        end
        def connect
            unless @connection
                logger.debug "Creating new connection to AWS"
                @connection = Fog::Compute.new({
                    :provider               => 'AWS',
                    :region                 => @region,
                    :aws_access_key_id      => @provider_access_user,
                    :aws_secret_access_key  => @provider_access_pass
                })
            end
            @connection
        end
        
        def all_named_servers(server_name)
            connect.servers.all({'tag:Name' => server_name})
        end

        def is_elastic_ip(ip_address)
            connect.addresses.each do |address|
                if address.public_ip == ip_address
                    return address
                end
            end
            return false
        end
    end
end