from common import *  # NOQA
from cattle import ApiError
import yaml


def _get_agent_for_container(context, container):
    agent = None
    for map in container.hosts()[0].instanceHostMaps():
        try:
            c = map.instance()
        except Exception:
            continue
        if c.agentId is not None:
            agent = c.agent()
            break

    if agent is None:
        client = context.client
        env = client.create_environment(name='env-' + random_str())
        svc = client.create_service(name='agentglobal' + random_str(),
                                    launchConfig={
            'imageUuid': context.image_uuid,
            "labels": {
                'io.rancher.scheduler.global': 'true',
                'io.rancher.container.create_agent': 'true'
            }
        }, stackId=env.id)
        svc = wait_state(client, svc, 'inactive')
        client.wait_success(svc.activate())

    for map in container.hosts()[0].instanceHostMaps():
        try:
            c = map.instance()
        except Exception:
            continue
        if c.agentId is not None:
            agent = c.agent()
            break

    assert agent is not None
    return agent


def _get_agent_client(agent):
    creds = agent.account().credentials()

    api_key = [x for x in creds if x.kind == 'agentApiKey'][0]
    assert len(api_key)
    return api_client(api_key.publicValue, api_key.secretValue)


def test_upgrade_with_health(client, context, super_client):
    env = client.create_environment(name='env-' + random_str())
    image_uuid = context.image_uuid
    launch_config = {"imageUuid": image_uuid}

    svc = client.create_service(name=random_str(),
                                stackId=env.id,
                                launchConfig=launch_config,
                                scale=1)
    svc = client.wait_success(svc)
    assert svc.state == "inactive"

    env.activateservices()
    svc = client.wait_success(svc, 120)
    assert svc.state == "active"

    # upgrade the service and
    new_launch_config = {"imageUuid": image_uuid,
                         'healthCheck': {
                             'port': 80,
                         }}
    strategy = {"launchConfig": new_launch_config,
                "intervalMillis": 100}
    svc.upgrade_action(inServiceStrategy=strategy)

    def wait_for_map_count(service):
        m = super_client. \
            list_serviceExposeMap(serviceId=service.id, state='active')
        return len(m) == 2

    def wait_for_managed_map_count(service):
        m = super_client. \
            list_serviceExposeMap(serviceId=service.id, state='active',
                                  upgrade=False)
        return len(m) == 1

    wait_for_condition(client, svc, wait_for_map_count)
    wait_for_condition(client, svc, wait_for_managed_map_count)

    m = super_client. \
        list_serviceExposeMap(serviceId=svc.id, state='active', upgrade=False)

    c = super_client.reload(m[0].instance())

    wait_for(lambda: super_client.reload(c).state == 'running')

    c = super_client.reload(m[0].instance())

    hci = find_one(c.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c)

    assert hcihm.healthState == 'initializing'
    assert c.healthState == 'initializing'

    # shouldn't become upgraded at this point
    try:
        wait_for(lambda: super_client.reload(svc).state == 'upgraded',
                 timeout=10)
    except Exception:
        pass

    _update_healthy(agent, hcihm, c, super_client)

    wait_for(lambda: super_client.reload(svc).healthState == 'healthy')
    wait_for(lambda: super_client.reload(svc).state == 'upgraded')
    super_client.reload(svc).remove()


def test_upgrade_start_first_with_health(client, context, super_client):
    env = client.create_environment(name='env-' + random_str())

    image_uuid = context.image_uuid
    launch_config = {"imageUuid": image_uuid}

    svc = client.create_service(name=random_str(),
                                stackId=env.id,
                                launchConfig=launch_config,
                                scale=1)
    svc = client.wait_success(svc)
    assert svc.state == "inactive"

    env.activateservices()
    svc = client.wait_success(svc, 120)
    assert svc.state == "active"

    # upgrade the service and
    # check that c3 and c2 got the same ip
    new_launch_config = {"imageUuid": image_uuid,
                         'healthCheck': {
                             'port': 80,
                         }}
    strategy = {"launchConfig": new_launch_config,
                "intervalMillis": 100,
                "startFirst": True}
    svc.upgrade_action(inServiceStrategy=strategy)

    def wait_for_map_count(service):
        m = super_client. \
            list_serviceExposeMap(serviceId=service.id, state='active')
        return len(m) == 2

    def wait_for_managed_map_count(service):
        m = super_client. \
            list_serviceExposeMap(serviceId=service.id, state='active',
                                  upgrade=False)
        return len(m) == 1

    wait_for_condition(client, svc, wait_for_map_count)
    wait_for_condition(client, svc, wait_for_managed_map_count)

    m = super_client. \
        list_serviceExposeMap(serviceId=svc.id, state='active', upgrade=False)

    c = super_client.reload(m[0].instance())

    wait_for(lambda: super_client.reload(c).state == 'running')

    c = super_client.reload(m[0].instance())

    hci = find_one(c.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c)

    assert hcihm.healthState == 'initializing'
    assert c.healthState == 'initializing'

    _update_healthy(agent, hcihm, c, super_client)

    wait_for(lambda: super_client.reload(svc).healthState == 'healthy')
    wait_for(lambda: super_client.reload(svc).state == 'upgraded')
    super_client.reload(svc).remove()


def test_health_check_create_instance(super_client, context):
    c = context.create_container(healthCheck={
        'port': 80,
    })

    assert c.healthCheck.port == 80

    c = super_client.reload(c)
    hci = find_one(c.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c)

    assert hcihm.healthState == 'initializing'

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)

    se = super_client.wait_success(se)
    assert se.state == 'created'
    assert se.accountId == c.accountId
    assert se.instanceId == c.id
    assert se.healthcheckInstanceId == hci.id

    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    assert hcihm.externalTimestamp == ts

    wait_for(lambda: super_client.reload(c).healthState == 'healthy')


def _create_svc_w_healthcheck(client, context):
    env = client.create_environment(name='env-' + random_str())
    return _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80,
        }
    }, stackId=env.id)


def test_health_check_create_service(super_client, context, client):
    service = _create_svc_w_healthcheck(client, context)

    maps = _wait_until_active_map_count(service, 1, client)
    expose_map = maps[0]
    c = super_client.reload(expose_map.instance())
    hci = find_one(c.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c)

    assert hcihm.healthState == 'initializing'
    assert c.healthState == 'initializing'

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='Something Bad',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'initializing'
    assert c.healthState == 'initializing'

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c).healthState == 'healthy')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c).healthState == 'healthy')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='INIT',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c).healthState == 'healthy')

    # restart the instance
    c = super_client.wait_success(c.restart())
    wait_for(lambda: super_client.reload(c).state == 'running')
    wait_for(lambda: super_client.reload(c).healthState == 'reinitializing')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='Something bad',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c).healthState == 'reinitializing')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c).healthState == 'healthy')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='INIT',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c).healthState == 'healthy')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c).healthState == 'healthy')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='Something Bad',
                                     healthcheckUuid=hcihm.uuid)

    se = super_client.wait_success(se)
    assert se.state == 'created'
    assert se.accountId == c.accountId
    assert se.instanceId == c.id
    assert se.healthcheckInstanceId == hci.id

    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'unhealthy'
    assert hcihm.externalTimestamp == ts

    wait_for(lambda: super_client.reload(c).healthState == 'unhealthy')
    wait_for(lambda: len(service.serviceExposeMaps()) == 1)
    remove_service(service)


def test_health_check_ip_retain(super_client, context, client):
    env = client.create_environment(name='env-' + random_str())
    service = _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80,
        }
    }, stackId=env.id, retainIp=True)

    maps = _wait_until_active_map_count(service, 1, client)
    expose_map = maps[0]
    c1 = super_client.reload(expose_map.instance())
    ip1 = c1.primaryIpAddress
    hci = find_one(c1.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c1)

    assert hcihm.healthState == 'initializing'
    assert c1.healthState == 'initializing'

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c1).healthState == 'healthy')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='Something Bad',
                                     healthcheckUuid=hcihm.uuid)

    se = super_client.wait_success(se)
    assert se.state == 'created'
    assert se.accountId == c1.accountId
    assert se.instanceId == c1.id
    assert se.healthcheckInstanceId == hci.id

    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'unhealthy'
    assert hcihm.externalTimestamp == ts

    wait_for(lambda: super_client.reload(c1).healthState == 'unhealthy')

    def check():
        s = super_client.reload(service)
        return s.instanceIds is not None and len(s.instanceIds) == 1 and \
            c1.id not in s.instanceIds

    wait_for(check)
    for e_map in service.serviceExposeMaps():
        if e_map.instance().id == c1.id:
            continue
        c2 = wait_state(super_client, e_map.instance(), 'running')
        assert c2.name == c1.name
        assert c2.primaryIpAddress == ip1
        break
    remove_service(service)


def test_health_state_stack(super_client, context, client):
    c = client
    env = client.create_environment(name='env-' + random_str())
    service = _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80,
        }
    }, stackId=env.id)

    i = 'initializing'
    wait_for(lambda: super_client.reload(service).healthState == i)
    wait_for(lambda: client.reload(env).healthState == i)

    maps = _wait_until_active_map_count(service, 1, client)
    expose_map = maps[0]
    c1 = super_client.reload(expose_map.instance())
    hci = find_one(c1.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c1)
    assert hcihm.healthState == 'initializing'
    assert c1.healthState == 'initializing'

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c1).healthState == 'healthy')

    wait_for(lambda: super_client.reload(service).healthState == 'healthy')
    wait_for(lambda: c.reload(env).healthState == 'healthy')

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='Something Bad',
                                     healthcheckUuid=hcihm.uuid)

    se = super_client.wait_success(se)
    assert se.state == 'created'
    assert se.accountId == c1.accountId
    assert se.instanceId == c1.id
    assert se.healthcheckInstanceId == hci.id

    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'unhealthy'
    assert hcihm.externalTimestamp == ts

    wait_for(lambda: super_client.reload(c1).healthState == 'unhealthy')


def test_health_state_start_once(super_client, context, client):
    c = client
    env = client.create_environment(name='env-' + random_str())
    labels = {"io.rancher.container.start_once": "true"}
    svc = _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80,
        },
        'labels': labels
    }, stackId=env.id)

    wait_for(lambda: super_client.reload(svc).healthState == 'initializing')
    wait_for(lambda: c.reload(env).healthState == 'initializing')

    maps = _wait_until_active_map_count(svc, 1, client)
    expose_map = maps[0]
    c1 = super_client.reload(expose_map.instance())
    hci = find_one(c1.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c1)
    assert hcihm.healthState == 'initializing'
    assert c1.healthState == 'initializing'

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c1).healthState == 'healthy')

    wait_for(lambda: super_client.reload(svc).healthState == 'healthy')
    wait_for(lambda: c.reload(env).healthState == 'healthy')

    super_client.wait_success(c1.stop())

    wait_for(lambda: super_client.reload(svc).healthState == 'started-once')
    wait_for(lambda: c.reload(env).healthState == 'started-once')
    remove_service(svc)


def test_health_state_sidekick_start_once(super_client, context, client):
    env = client.create_environment(name='env-' + random_str())
    labels = {"io.rancher.container.start_once": "true"}
    slc = {'imageUuid': context.image_uuid, 'name': "test1"}
    svc = _create_and_activate_svc(client, 2, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'labels': labels
    }, stackId=env.id, secondaryLaunchConfigs=[slc])

    wait_for(lambda: super_client.reload(svc).healthState == 'healthy')
    wait_for(lambda: client.reload(env).healthState == 'healthy')
    remove_service(svc)


def test_health_state_selectors(context, client):
    env = client.create_environment(name='env-' + random_str())
    labels = {'foo': "bar"}
    container1 = client.create_container(imageUuid=context.image_uuid,
                                         startOnCreate=True,
                                         labels=labels)
    container1 = client.wait_success(container1)
    assert container1.state == "running"

    launch_config = {"imageUuid": "rancher/none"}
    service = _create_and_activate_svc(client, 1, name=random_str(),
                                       stackId=env.id,
                                       launchConfig=launch_config,
                                       selectorContainer="foo=bar")
    service = client.wait_success(service)
    assert service.selectorContainer == "foo=bar"

    wait_for(lambda: client.reload(service).healthState == 'healthy')
    wait_for(lambda: client.reload(env).healthState == 'healthy')
    remove_service(service)


def test_svc_health_state(context, client):
    env = client.create_environment(name='env-' + random_str())

    launch_config = {"imageUuid": context.image_uuid}
    service = _create_and_activate_svc(client, 1, name=random_str(),
                                       stackId=env.id,
                                       launchConfig=launch_config)

    wait_for(lambda: client.reload(service).healthState == 'healthy')
    wait_for(lambda: client.reload(env).healthState == 'healthy')
    remove_service(service)


def test_health_check_init_timeout(super_client, context, client):
    env = client.create_environment(name='env-' + random_str())
    service = _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80,
            'initializingTimeout': 2000,
        }
    }, stackId=env.id)
    h_c = service.launchConfig.healthCheck
    assert h_c.initializingTimeout == 2000

    maps = _wait_until_active_map_count(service, 1, client)
    expose_map = maps[0]
    c = super_client.reload(expose_map.instance())
    hci = find_one(c.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)

    assert hcihm.healthState == 'initializing'
    assert c.healthState == 'initializing'

    # wait for the instance to be removed
    wait_for_condition(client, c, lambda x: x.removed is None)
    remove_service(service)


def test_health_check_reinit_timeout(super_client, context, client):
    env = client.create_environment(name='env-' + random_str())
    service = _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80,
            'reinitializingTimeout': 1,
        }
    }, stackId=env.id)
    h_c = service.launchConfig.healthCheck
    assert h_c.reinitializingTimeout == 1

    maps = _wait_until_active_map_count(service, 1, client)
    expose_map = maps[0]
    c = super_client.reload(expose_map.instance())
    hci = find_one(c.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c)

    assert hcihm.healthState == 'initializing'
    assert c.healthState == 'initializing'

    ts = int(time.time())
    client = _get_agent_client(agent)
    se = client.create_service_event(externalTimestamp=ts,
                                     reportedHealth='UP',
                                     healthcheckUuid=hcihm.uuid)
    super_client.wait_success(se)
    hcihm = super_client.wait_success(super_client.reload(hcihm))
    assert hcihm.healthState == 'healthy'
    wait_for(lambda: super_client.reload(c).healthState == 'healthy')

    # restart the instance
    c = super_client.wait_success(c.restart())
    wait_for(lambda: super_client.reload(c).state == 'running')
    wait_for(lambda: super_client.reload(c).healthState == 'reinitializing')

    # wait for the instance to be removed
    wait_for_condition(super_client, c,
                       lambda x: x.removed is not None)
    remove_service(service)


def test_health_check_bad_external_timestamp(super_client, context, client):
    env = client.create_environment(name='env-' + random_str())
    service = _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80,
        }
    }, stackId=env.id)

    maps = _wait_until_active_map_count(service, 1, client)
    expose_map = maps[0]
    container = super_client.reload(expose_map.instance())
    hci = find_one(container.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, container)
    agent_client = _get_agent_client(agent)

    assert hcihm.healthState == 'initializing'

    with pytest.raises(ApiError) as e:
        agent_client.create_service_event(reportedHealth='Something Bad',
                                          healthcheckUuid=hcihm.uuid)
    assert e.value.error.code == 'MissingRequired'
    assert e.value.error.fieldName == 'externalTimestamp'
    remove_service(service)


def test_health_check_noop(super_client, context, client):
    env = client.create_environment(name='env-' + random_str())
    svc = _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80,
            'strategy': 'none'
        }
    }, stackId=env.id)

    assert svc.launchConfig.healthCheck.strategy == 'none'

    maps = _wait_until_active_map_count(svc, 1, client)
    expose_map = maps[0]
    c = super_client.reload(expose_map.instance())
    hci = find_one(c.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c)

    assert hcihm.healthState == 'initializing'
    assert c.healthState == 'initializing'

    hcihm = _update_healthy(agent, hcihm, c, super_client)
    svc = wait_state(client, svc, 'active')

    _update_unhealthy(agent, hcihm, c, super_client)

    svc = wait_state(client, svc, 'active')

    assert len(svc.serviceExposeMaps()) == 1
    c = super_client.wait_success(c)
    assert c.state == 'running'
    remove_service(svc)


def remove_service(svc):
    for i in range(1, 3):
        try:
            svc.remove()
            return
        except Exception:
            pass
        time.sleep(1)


def test_health_check_default(super_client, context, client):
    env = client.create_environment(name='env-' + random_str())
    svc = _create_and_activate_svc(client, 1, name='test', launchConfig={
        'imageUuid': context.image_uuid,
        'healthCheck': {
            'port': 80
        }
    }, stackId=env.id)

    expose_maps = svc.serviceExposeMaps()
    c1 = super_client.reload(expose_maps[0].instance())
    hci = find_one(c1.healthcheckInstances)
    hcihm = find_one(hci.healthcheckInstanceHostMaps)
    agent = _get_agent_for_container(context, c1)

    assert hcihm.healthState == 'initializing'
    assert c1.healthState == 'initializing'