# Copyright 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.

"""Integration tests for Instance API.
"""

import datetime
import pytest
import time
import logging

from acktest.resources import random_suffix_name
from acktest.k8s import resource as k8s
from e2e import service_marker, CRD_GROUP, CRD_VERSION, load_ec2_resource
from e2e.replacement_values import REPLACEMENT_VALUES
from e2e.bootstrap_resources import get_bootstrap_resources

RESOURCE_PLURAL = "instances"
# highly available instance type for deterministic testing
INSTANCE_TYPE = "m4.large"
INSTANCE_AMI = "Amazon Linux 2 Kernel"
INSTANCE_TAG_KEY = "owner"
INSTANCE_TAG_VAL = "ack-controller"

CREATE_WAIT_AFTER_SECONDS = 10
DELETE_WAIT_AFTER_SECONDS = 10
TIMEOUT_SECONDS = 300

def get_instance(ec2_client, instance_id: str) -> dict:
    instance = None
    try:
        resp = ec2_client.describe_instances(
            InstanceIds=[instance_id]
        )
        instance = resp["Reservations"][0]["Instances"][0]
    except Exception as e:
        logging.debug(e)
    finally:
        return instance

def get_instance_state(ec2_client, instance_id):
    instance_state = None
    try:
        instance = get_instance(ec2_client, instance_id)
        instance_state = instance["State"]["Name"]
    except Exception as e:
        logging.debug(e)
    finally:
        return instance_state

def wait_for_instance_or_die(ec2_client, instance_id, desired_state, timeout_sec):
    while True:
        now = datetime.datetime.now()
        timeout = now + datetime.timedelta(seconds=timeout_sec)
        if datetime.datetime.now() >= timeout:
            pytest.fail(f"Timed out waiting for Instance to enter {desired_state} state")
        time.sleep(DELETE_WAIT_AFTER_SECONDS)
        instance_state = get_instance_state(ec2_client, instance_id)
        if instance_state == desired_state:
            break

def get_ami_id(ec2_client):
    try:
        # Use latest AL2
        resp = ec2_client.describe_images(
            Owners=['amazon'],
            Filters=[
                {"Name": "architecture", "Values": ['x86_64']},
                {"Name": "state", "Values": ['available']},
                {"Name": "virtualization-type", "Values": ['hvm']},
                ],
        )
        for image in resp['Images']:
            if 'Description' in image:
                if INSTANCE_AMI in image['Description']:
                    return image['ImageId']
    except Exception as e:
        logging.debug(e)


@pytest.fixture
def instance(ec2_client):
    test_resource_values = REPLACEMENT_VALUES.copy()
    resource_name = random_suffix_name("instance-ack-test", 24)
    test_vpc = get_bootstrap_resources().SharedTestVPC
    subnet_id = test_vpc.public_subnets.subnet_ids[0]
        
    ami_id = get_ami_id(ec2_client)
    test_resource_values["INSTANCE_NAME"] = resource_name
    test_resource_values["INSTANCE_AMI_ID"] = ami_id
    test_resource_values["INSTANCE_TYPE"] = INSTANCE_TYPE
    test_resource_values["INSTANCE_SUBNET_ID"] = subnet_id
    test_resource_values["INSTANCE_TAG_KEY"] = INSTANCE_TAG_KEY
    test_resource_values["INSTANCE_TAG_VAL"] = INSTANCE_TAG_VAL

    # Load Instance CR
    resource_data = load_ec2_resource(
        "instance",
        additional_replacements=test_resource_values,
    )
    logging.debug(resource_data)

    # Create k8s resource
    ref = k8s.CustomResourceReference(
        CRD_GROUP, CRD_VERSION, RESOURCE_PLURAL,
        resource_name, namespace="default",
    )
    k8s.create_custom_resource(ref, resource_data)
    cr = k8s.wait_resource_consumed_by_controller(ref)

    assert cr is not None
    assert k8s.get_resource_exists(ref)
    
    yield (ref, cr)

    # Delete the instance when tests complete
    try:
        _, deleted = k8s.delete_custom_resource(ref, 3, 10)
        assert deleted
    except:
        pass

@service_marker
@pytest.mark.canary
class TestInstance:
    def test_create_delete(self, ec2_client, instance):
        (ref, cr) = instance
        resource_id = cr["status"]["instanceID"]

        time.sleep(CREATE_WAIT_AFTER_SECONDS)

        # Check Instance exists
        instance = get_instance(ec2_client, resource_id)
        assert instance is not None
        
        # Give time for instance to come up
        wait_for_instance_or_die(ec2_client, resource_id, 'running', TIMEOUT_SECONDS)

        # Validate instance tags
        instance_tags = instance["Tags"]
        tag_present = False
        for t in instance_tags:
            if (t['Key'] == INSTANCE_TAG_KEY and
                    t['Value'] == INSTANCE_TAG_VAL):
                tag_present = True
        assert tag_present

        # Delete k8s resource
        _, deleted = k8s.delete_custom_resource(ref, 2, 5)
        assert deleted is True

        # Reservation still exists, but instance will commence termination
        # State needs to be 'terminated' in order to remove the dependency on the shared subnet
        # for successful test cleanup
        wait_for_instance_or_die(ec2_client, resource_id, 'terminated', TIMEOUT_SECONDS)
[((93, 27, 93, 52), 'e2e.replacement_values.REPLACEMENT_VALUES.copy', 'REPLACEMENT_VALUES.copy', ({}, {}), '()', False, 'from e2e.replacement_values import REPLACEMENT_VALUES\n'), ((94, 20, 94, 63), 'acktest.resources.random_suffix_name', 'random_suffix_name', ({(94, 39, 94, 58): '"""instance-ack-test"""', (94, 60, 94, 62): '24'}, {}), "('instance-ack-test', 24)", False, 'from acktest.resources import random_suffix_name\n'), ((107, 20, 110, 5), 'e2e.load_ec2_resource', 'load_ec2_resource', (), '', False, 'from e2e import service_marker, CRD_GROUP, CRD_VERSION, load_ec2_resource\n'), ((111, 4, 111, 32), 'logging.debug', 'logging.debug', ({(111, 18, 111, 31): 'resource_data'}, {}), '(resource_data)', False, 'import logging\n'), ((114, 10, 117, 5), 'acktest.k8s.resource.CustomResourceReference', 'k8s.CustomResourceReference', (), '', True, 'from acktest.k8s import resource as k8s\n'), ((118, 4, 118, 50), 'acktest.k8s.resource.create_custom_resource', 'k8s.create_custom_resource', ({(118, 31, 118, 34): 'ref', (118, 36, 118, 49): 'resource_data'}, {}), '(ref, resource_data)', True, 'from acktest.k8s import resource as k8s\n'), ((119, 9, 119, 54), 'acktest.k8s.resource.wait_resource_consumed_by_controller', 'k8s.wait_resource_consumed_by_controller', ({(119, 50, 119, 53): 'ref'}, {}), '(ref)', True, 'from acktest.k8s import resource as k8s\n'), ((122, 11, 122, 39), 'acktest.k8s.resource.get_resource_exists', 'k8s.get_resource_exists', ({(122, 35, 122, 38): 'ref'}, {}), '(ref)', True, 'from acktest.k8s import resource as k8s\n'), ((63, 14, 63, 37), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((67, 8, 67, 45), 'time.sleep', 'time.sleep', ({(67, 19, 67, 44): 'DELETE_WAIT_AFTER_SECONDS'}, {}), '(DELETE_WAIT_AFTER_SECONDS)', False, 'import time\n'), ((95, 15, 95, 40), 'e2e.bootstrap_resources.get_bootstrap_resources', 'get_bootstrap_resources', ({}, {}), '()', False, 'from e2e.bootstrap_resources import get_bootstrap_resources\n'), ((128, 21, 128, 59), 'acktest.k8s.resource.delete_custom_resource', 'k8s.delete_custom_resource', ({(128, 48, 128, 51): 'ref', (128, 53, 128, 54): '3', (128, 56, 128, 58): '10'}, {}), '(ref, 3, 10)', True, 'from acktest.k8s import resource as k8s\n'), ((140, 8, 140, 45), 'time.sleep', 'time.sleep', ({(140, 19, 140, 44): 'CREATE_WAIT_AFTER_SECONDS'}, {}), '(CREATE_WAIT_AFTER_SECONDS)', False, 'import time\n'), ((159, 21, 159, 58), 'acktest.k8s.resource.delete_custom_resource', 'k8s.delete_custom_resource', ({(159, 48, 159, 51): 'ref', (159, 53, 159, 54): '2', (159, 56, 159, 57): '5'}, {}), '(ref, 2, 5)', True, 'from acktest.k8s import resource as k8s\n'), ((47, 8, 47, 24), 'logging.debug', 'logging.debug', ({(47, 22, 47, 23): 'e'}, {}), '(e)', False, 'import logging\n'), ((57, 8, 57, 24), 'logging.debug', 'logging.debug', ({(57, 22, 57, 23): 'e'}, {}), '(e)', False, 'import logging\n'), ((64, 24, 64, 63), 'datetime.timedelta', 'datetime.timedelta', (), '', False, 'import datetime\n'), ((65, 11, 65, 34), 'datetime.datetime.now', 'datetime.datetime.now', ({}, {}), '()', False, 'import datetime\n'), ((66, 12, 66, 89), 'pytest.fail', 'pytest.fail', ({(66, 24, 66, 88): 'f"""Timed out waiting for Instance to enter {desired_state} state"""'}, {}), "(f'Timed out waiting for Instance to enter {desired_state} state')", False, 'import pytest\n'), ((88, 8, 88, 24), 'logging.debug', 'logging.debug', ({(88, 22, 88, 23): 'e'}, {}), '(e)', False, 'import logging\n')]