#! /usr/bin/env python
# -*- coding:utf-8 -*-
import shutil
import json


def main():
   quit_flag = False
   while not quit_flag:
       item_list()
       choice = input("make your choice:").strip()
       if choice == "1":
           backend_info = input("input your fetch backend:").strip()
           result = fetch(backend_info)
           print(result)
       elif choice == "2":
           back_up()
           backend_record_dict = input_json()  # 要求输入json格式
           for key in backend_record_dict:
               backend = key
               record = backend_record_dict[key]
           delete(backend,record)
       elif choice == "3":
           back_up()
           backend_record_dict = input_json()  # 要求输入json格式
           for key in backend_record_dict:
               backend = key
               record = backend_record_dict[key]
           add(backend,record)
       elif choice == "4":
           back_up()
           backend_record_dict = input_json()  # 要求输入json格式
           for key in backend_record_dict:
               backend = key
               record = backend_record_dict[key]
           modify(backend,record)
       elif choice == "q" or choice == "quit":
           quit_flag = True
           exit("quit")
       else:
           print("invalid type")


def item_list():
    lis = ["1.查询", "2.删除", "3.添加", "4.修改"]
    print(lis)


def back_up():
    shutil.copyfile('haproxy', 'haproxy.bak')


def fetch(backend):
    flag = False
    fetch_list = []
    with open("haproxy", "r", encoding="utf-8") as f:
        for line in f:
            if line.strip() == "backend %s" %(backend):
                flag = True  # 读取到backend开头的信息，标记为真
                continue
            if line.strip().startswith("backend"):  # 下一个backend开头
                flag = False
            if flag and line.strip():
                fetch_list.append(line.strip())
    return fetch_list


def fetch1(backend):
    # 打印bankend和server   防止后面的判断不准确
    flag = False
    fetch_list = []
    with open("haproxy", "r", encoding="utf-8") as f:
        for line in f:
            if line.strip() == "backend %s" %(backend):
                fetch_list.append(backend)
                flag = True  # 读取到backend开头的信息，标记为真
                continue
            if line.strip().startswith("backend"):  # 下一个backend开头
                flag = False
            if flag and line.strip():
                fetch_list.append(line.strip())
    return fetch_list


def writer(backend, record_list):
    # 保留backend，删除或增加server
    with open("haproxy", "r", encoding="utf-8") as old, open("new.cfg", "w", encoding="utf-8") as new:
        flag = False
        for line in old:
            if line.strip().startswith('backend') and line.strip() == "backend %s" % backend:
                flag = True
                new.write(line)
                for new_line in record_list:
                    new.write(" " * 4 + new_line + "\n")
                continue  # 跳到下一次循环，避免下一个backend写二次
            if flag and line.strip().startswith("backend"):  # 下一个backend
                flag = False
                new.write(line)
                continue  # 退出此循环，避免server信息二次写入
            if line.strip() and not flag:
                new.write(line)


def writer1(backend, record_list):
    # 删除backend，删除或增加server
    with open("haproxy", "r", encoding="utf-8") as old, open("new.cfg", "w", encoding="utf-8") as new:
        flag = False
        for line in old:
            if line.strip().startswith('backend') and line.strip() == "backend %s" % backend:
                flag = True
                #new.write(line)
                for new_line in record_list:
                    new.write(" " * 4 + new_line + "\n")
                    new.flush()
                continue  # 跳到下一次循环，避免下一个backend写二次
            if flag and line.strip().startswith("backend"):  # 下一个backend
                flag = False
                new.write(line)
                new.flush()
                continue  # 退出此循环，避免server信息二次写入
            if line.strip() and not flag:
                new.write(line)
                new.flush()


def writer2(backend, record_list):
    # 只删除backend
    with open("haproxy", "r", encoding="utf-8") as old, open("new.cfg", "w", encoding="utf-8") as new:
        flag = False
        for line in old:
            if line.strip().startswith('backend') and line.strip() == "backend %s" % backend:
                flag = True
                #new.write(line)
                # for new_line in record_list:
                #     new.write(" " * 4 + new_line + "\n")
                continue  # 跳到下一次循环，避免下一个backend写二次
            if flag and line.strip().startswith("backend"):  # 下一个backend
                flag = False
                new.write(line)
                new.flush()
                continue  # 退出此循环，避免server信息二次写入
            if line.strip() and not flag:
                new.write(line)
                new.flush()


def add(backend, record):
    global change_flag
    record_list = fetch(backend)  # 查找是否存在记录
    if not record_list:  # backend不存在, record不存在
        with open('haproxy', 'r',encoding="utf-8") as old, open('new.cfg', 'w', encoding="utf-8") as new:
            for line in old:
                new.write(line)
            # 添加新记录
            new.write("\nbackend " + backend + "\n")
            new.write(" " * 4 + record + "\n")
        print("\033[32;1mAdd done\033[0m")
        change_flag = True
    else:  # backend存在，record存在
        if record in record_list:
            print("\033[31;1mRecord already in it,Nothing to do!\033[0m")
            change_flag = False
        else:  # backend存在，record不存在
            record_list.append(record)
            writer(backend, record_list)
            print("\033[32;1mAdd done\033[0m")
            change_flag = True
    if change_flag is True:  # 文件有修改，才进行文件更新
        # 将操作结果生效
        shutil.move("new.cfg", "haproxy")


def delete(backend, record):
    global change_flag
    record_list = fetch1(backend)  # 查找是否存在记录
    print(record_list)
    if not record_list:  # backend不存在, record不存在
        print("Not match the backend,no need delete!".center(50, "#"))
        change_flag = False
    else:  # backend存在，record存在
        if record in record_list:
            record_list.remove(record)  # 移除元素
            writer1(backend, record_list)  # 写入   删除backend和server
            print("\033[31;1mDelete done\033[0m")
            change_flag = True
        else:  # backend存在，record不存在  需要删除
            # print("Only match backend,no need delete!".center(50, "#"))
            # change_flag = False
            writer2(backend,record)
            print("\033[31;1mDelete done\033[0m")
            change_flag = True
    if change_flag is True:  # 文件有修改，才进行文件更新
        # 将操作结果生效
        shutil.move("new.cfg", "haproxy")
    return change_flag


def modify(backend, record):
    global change_flag
    record_list = fetch1(backend)  # 查找是否存在记录
    if not record_list:  # backend不存在, record不存在
        with open('haproxy', 'r',encoding="utf-8") as old, open('new.cfg', 'w', encoding="utf-8") as new:
            for line in old:
                new.write(line)
            # 添加新记录
            new.write("\nbackend " + backend + "\n")
            new.write(" " * 4 + record + "\n")
        change_flag = True
        if change_flag is True:  # 文件有修改，才进行文件更新
            # 将操作结果生效
            shutil.move("new.cfg", "haproxy")
    else:  # backend存在或者record存在
        delete(backend, record)
        add(backend, record)
    return change_flag


def input_json():
    # 判断输入，要求为json格式
    continue_flag = False
    while continue_flag is not True:
        backend_record = input("Input backend info(json):").strip()
        try:
            backend_record_dict = json.loads(backend_record)
        except Exception as e:
            print("\033[31;1mInput not a json data type!\033[0m")
            continue
        continue_flag = True
    return backend_record_dict

if __name__ == '__main__':
    main()