from argparse import ArgumentParser
from socket import socket, AF_INET, SOCK_DGRAM, IPPROTO_IP, IP_TOS
from time import sleep
import re
import time
import sys
from loguru import logger


def parse_bandwidth(input_str):
    pattern = r'^(\d+(?:\.\d+)?)\s*([kKmMgGtT])bps$'
    match = re.match(pattern, input_str)
    
    if match:
        value = float(match.group(1))
        unit = match.group(2).lower()
        
        multipliers = {
            'k': 1e3,
            'm': 1e6,
            'g': 1e9,
            't': 1e12
        }
        
        multiplier = multipliers.get(unit, 1.0)
        return value * multiplier
    else:
        raise ValueError(f"Invalid bandwidth: {input_str}")


def main():
    parser = ArgumentParser()
    parser.add_argument('--host', type=str, default='localhost')
    parser.add_argument('--port', type=int, default=12345)
    parser.add_argument("--log-level", type=str, default="INFO")
    parser.add_argument("--bandwidth", type=str, default="1Mbps")
    parser.add_argument("--duration", type=int, default=10)
    parser.add_argument("--dscp", action='append', default=[])
    args = parser.parse_args()
    logger.remove()
    logger.add(sys.stderr, level=args.log_level)
    PACKET_SIZE = 1280
    bandwidth = parse_bandwidth(args.bandwidth)
    time_between_packets = PACKET_SIZE / bandwidth
    packet_cnt = 0
    logger.info(f"Connecting to {args.host}:{args.port}")
    

    sockets = []
    dscp = [int(x) for x in args.dscp]
    if dscp:
        for value in dscp:
            sock = socket(AF_INET, SOCK_DGRAM)
            sock.setsockopt(IPPROTO_IP, IP_TOS, value << 2)
            sockets.append(sock)
    else:
        sock = socket(AF_INET, SOCK_DGRAM)
        sockets.append(sock)
    CODE_WORD = "CAFEABBA"
    payload = (CODE_WORD + "A" * (PACKET_SIZE - len(CODE_WORD))).encode()
    t0 = time.time()
    while True:
        logger.debug("Sending message:")
        for sock in sockets:
            sock.sendto(payload, (args.host, args.port))
        packet_cnt += 1
        t1 = time.time() - t0
        if t1 > args.duration:
            break
        time_to_sleep = time_between_packets * packet_cnt - t1
        if time_to_sleep > 0:
            sleep(time_between_packets * packet_cnt - t1)
        else:
            logger.warning("We are too slow: {}s", -time_to_sleep)
        
if __name__ == '__main__':
    main()