#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################################
# Copyright (c) 2020 Huawei Technologies Co.,Ltd.
#
# openGauss is licensed under Mulan PSL v2.
# You can use this software according to the terms
# and conditions of the Mulan PSL v2.
# You may obtain a copy of Mulan PSL v2 at:
#
#          http://license.coscl.org.cn/MulanPSL2
#
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS,
# WITHOUT WARRANTIES OF ANY KIND,
# EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
# MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.
# ----------------------------------------------------------------------------
import socket
import time
import sys
import ipaddress

listen_ip = "localhost"
listen_port = 31111
run_mode = 0  # 0:connect, 1:send, 2:recv
NET_IPV4 = "ipv4"
NET_IPV6 = "ipv6"

def get_sockects(ip_address):
    try:
        ip = ipaddress.ip_address(ip_address)
        if ip.version == 6:
            sockets = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)
            return sockets
        else:
            sockets = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            return sockets
    except ValueError:
        print("ip_address:%s is unvalid ip.")

def send_main():
    try:
        global listen_ip
        global listen_port
        buf = "this is a test !" * 512 # buf  8192  block
        sockets = get_sockects(listen_ip)
        sockets.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
        print(listen_ip+":"+listen_port)
        while(sockets.connect_ex((listen_ip, int(listen_port))) != 0):
            print("connect failed:%m\n")
            time.sleep(1)
        print("connect succeed, dest[%s:%d], mode[%s]\n", listen_ip, listen_port, "tcp")
        print("send satrt, dest[%s:%d], mode[%s]\n", listen_ip, listen_port, "tcp")
        i = 0
        while True:
            i = i + 1
            n = sockets.send(buf.encode())
            if n == 0:
                print("send failed:%m\n")
                break
            print("%d send:%s, len=%d\n", i, buf, n)
    except Exception as e:
        print(str(e))

def recv_main():
    try:
        global listen_ip
        global listen_port
        sockets = get_sockects(listen_ip)
        sockets.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, True)
        sockets.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
        sockets.bind((listen_ip, int(listen_port)))
        sockets.listen(128)
        while True:
            client, addr = sockets.accept()
            print('client:', client)
            print('addr:', addr)
            while True:
                data = client.recv(8192)
                print(data.decode())
                if not data:
                    client.close()
                    break
    except Exception as e:
        print(str(e))
def connect_main():
    sockets = get_sockects(listen_ip)
    sockets.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
    if sockets.connect_ex((listen_ip, int(listen_port))) != 0:
        print("Failed to connect %s:%d on %s mode:%m.\n",
              listen_ip, listen_port, "tcp")
    else:
        print("Succeed to connect %s:%d on %s mode.\n",
              listen_ip, listen_port, "tcp")

def invalid_argument():
    print("usage: ./speed_test recv/send/connect recv_ip "
          "recv_port [sctp] [msg_len] [debug]\n")
    print("example: ./speed_test recv 127.0.0.1 10001\n")
    exit(1)

if __name__ == '__main__':
    if len(sys.argv) < 5:
       invalid_argument()
    if sys.argv[1] == "send":
        run_mode = 1
    elif sys.argv[1] == "recv":
        run_mode = 2
    listen_ip = sys.argv[2]
    listen_port = sys.argv[3]
    if run_mode == 1:
        send_main()
    elif run_mode == 2:
        recv_main()
    else:
        connect_main()