#!/bin/env python
# coding=utf-8
import logging
import time
import docker
import sys
import requests
import json
import queue
import threading
from aliyunsdkcore import client
from aliyunsdkecs.request.v20140526.DescribeNetworkInterfacesRequest import DescribeNetworkInterfacesRequest
from aliyunsdkcore.request import CommonRequest
from aliyunsdkslb.request.v20140515.DescribeLoadBalancerAttributeRequest import DescribeLoadBalancerAttributeRequest
from concurrent.futures import ThreadPoolExecutor

logging.basicConfig(level=logging.INFO,
                    format='%(asctime)s %(filename)s[line:%(lineno)d] %(levelname)s %(message)s',
                    datefmt='%a, %d %b %Y %H:%M:%S')

# ##### 修改负载均衡权限
AccessKey = 'LTAIuPbTPL3LDDKN'
AccessSecret = 'ORcNedKwWuwVtcq4IRFtUDZgS0b1le'
RegionId = 'cn-hangzhou'
clt = client.AcsClient(AccessKey, AccessSecret, RegionId)

slb_id_1 = 'lb-bp1werfophtsjzfr76njm'


def getInstanceId(slb_id):
    """获取slb下所有服务器信息"""
    request = DescribeLoadBalancerAttributeRequest()
    request.set_accept_format('json')
    request.set_LoadBalancerId(slb_id)
    response = clt.do_action_with_exception(request)
    return json.loads(response)


def setWeight(slb_id, instance_id, weight):
    """
    设置slb权重
    :param slb_id: slb_id
    :param instance_id: 服务器id
    :param weight: 权重值
    :return: None
    """
    BackendServers = [{"ServerId": instance_id, "Weight": weight}]
    request = CommonRequest()
    request.set_accept_format('json')
    request.set_domain('slb.aliyuncs.com')
    request.set_version('2014-05-15')
    request.set_method('POST')
    request.set_action_name('SetBackendServers')
    request.add_query_param('BackendServers', BackendServers)
    request.add_query_param('LoadBalancerId', slb_id)
    try:
        response = clt.do_action_with_exception(request)
    except Exception as e:
        logging.error(e)


def set_weight_for_more(slb_id, instance_id_list, weight):
    """
    同时设置多台服务器的slb权重，权重一样
    :param slb_id: slb_id
    :param instance_id_list: 服务器id list
    :param weight: 权重值
    :return: None
    """
    BackendServers = [{"ServerId": instance_id, "Weight": weight} for instance_id in instance_id_list]
    request = CommonRequest()
    request.set_accept_format('json')
    request.set_domain('slb.aliyuncs.com')
    request.set_version('2014-05-15')
    request.set_method('POST')
    request.set_action_name('SetBackendServers')
    request.add_query_param('BackendServers', BackendServers)
    request.add_query_param('LoadBalancerId', slb_id)
    try:
        response = clt.do_action_with_exception(request)
    except Exception as e:
        logging.error(e)


def getIpadd(instance_id):
    """获取实例IP地址"""
    request = DescribeNetworkInterfacesRequest()
    request.set_accept_format('json')
    request.set_InstanceId(instance_id)
    response = clt.do_action_with_exception(request)
    request_content = json.loads(response)
    IpAddr = request_content['NetworkInterfaceSets']['NetworkInterfaceSet'][0]['PrivateIpAddress']
    return IpAddr


def checkHealth(ipadd):
    """健康检查"""
    while True:
        health_url = 'http://%s:5001/healthcheck' %(ipadd)
        header = {"Content-Type": "application/json"}
        try:
            health_code = requests.get(health_url).status_code
        except Exception as e:
           continue
        if health_code == 200:
            return False


def setInstanceWeightProcess(instance_id):
    """服务更新完之后逐步修改服务器的权重值"""
    # 直接加载100会出现502。权重值每次增加20，每5s修改一次
    for i in range(1, 6):
        weight = i * 20
        setWeight(slb_id_1, instance_id, weight)
        time.sleep(1)


def restart():
    """重启服务"""
    # 等待10s,处理剩余请求
    time.sleep(10)
    global success_count, finished_instance_id_list
    apps = 'rov-server'
    while not q2.empty():
        instance_id = q2.get()
        ipadd = getIpadd(instance_id)
        logging.info(f"服务器信息：{instance_id}/{ipadd}")
        client = docker.DockerClient(base_url=f'tcp://{ipadd}:2375', timeout=60)
        try:
            # 重启容器
            id = client.containers.get(apps)
            id.restart()
            checkHealth(ipadd)
            print(f"{ipadd}: 权重修改中......")
            # 修改权重
            setInstanceWeightProcess(instance_id)
            success_count = success_count + 1
            finished_instance_id_list.append(instance_id)
            print(f"重启进度: {success_count}/{total}")
        except Exception as e:
            raise e
            sys.exit()  # 容器启动失败立即退出更新


class MyThread(threading.Thread):
    def __init__(self, func):
        threading.Thread.__init__(self)
        self.func = func

    def run(self):
        self.func()


if __name__ == '__main__':
    # 更新完成计数
    success_count = 0
    threads = []
    # 线程数
    thread_num = 2
    # 获取slb下所有服务器信息
    res = getInstanceId(slb_id_1)
    # slb下服务器总数
    total = len(res["BackendServers"]["BackendServer"])
    q1 = queue.Queue()
    q2 = queue.Queue()
    # 获取slb下所有服务器的 instance_id
    if res["BackendServers"]["BackendServer"]:
        for i in range((len(res["BackendServers"]["BackendServer"]))):
            instance_id = res["BackendServers"]["BackendServer"][i]["ServerId"]
            q1.put(instance_id)

    while not q1.empty():
        # 摘流量
        for i in range(thread_num):
            if q1.empty():
                break
            instance_id = q1.get()
            # 设置slb权重为0
            setWeight(slb_id_1, instance_id, 0)
            q2.put(instance_id)
            time.sleep(1)

        # 挂流量
        finished_instance_id_list = []
        for i in range(thread_num):
            # 开启线程
            thread = MyThread(restart)
            thread.start()
            threads.append(thread)
        for thread in threads:
            thread.join()
        # 权重补充（一秒内同时请求两个服务器修改权重，可能会有一个失败）
        print(finished_instance_id_list)
        set_weight_for_more(slb_id=slb_id_1, instance_id_list=finished_instance_id_list, weight=100)
