import time
import threading
import argparse
import mooncake.engine as engine

class AtomicCounter:
    def __init__(self, initial=0):
        self._value = initial
        self._lock = threading.Lock()

    def inc(self, num=1):
        with self._lock:
            self._value += num
            return self._value

    def dec(self, num=1):
        with self._lock:
            self._value -= num
            return self._value

    def get_and_reset(self):
        with self._lock:
            r = self._value
            self._value = 0
            return r

counter = AtomicCounter()

# Global variable to store data size
data_size = 1024 * 1024  # Default 1MB
test_data = None

def print_qps():
    while True:
        time.sleep(1)
        val = counter.get_and_reset()
        if val == 0:
            continue
        print("bandwidth:", 8 * val * data_size / (1024*1024*1024), "GB/s")

def send_data(client, target_url):
    while True:
        try:
            result = client.send_data(target_url, test_data)
            if result < 0:
                print(f"Warning: send_data returned error code {result}")
                time.sleep(0.01)  # Brief delay on error to avoid tight loop
            else:
                counter.inc()
        except Exception as e:
            print(f"Error sending data to {target_url}: {e}")
            time.sleep(0.1)  # Delay on exception to avoid rapid retry

def run_server(bind_url, data_size_mb=1):
    """Run server mode"""
    global data_size, test_data
    data_size = data_size_mb * 1024 * 1024
    test_data = b'\x00' * data_size
    
    print(f"Starting server on {bind_url} with {data_size_mb}MB data packets")
    
    RpcInterface = engine.RpcInterface
    server = RpcInterface()
    server.initialize_server(bind_url, thread_count=8)
    server.start_server_async()
    
    # Start QPS statistics thread
    thread = threading.Thread(target=print_qps)
    thread.daemon = True
    thread.start()
    
    print(f"Server started on {bind_url}, press Ctrl+C to stop")
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\nServer stopping...")
        server.stop_server()

def run_client(target_url, num_threads=8, data_size_mb=1):
    """Run client mode"""
    global data_size, test_data
    data_size = data_size_mb * 1024 * 1024
    test_data = b'\x00' * data_size
    
    print(f"Starting client, connecting to {target_url} with {num_threads} threads, {data_size_mb}MB data packets")
    
    RpcInterface = engine.RpcInterface
    client = RpcInterface()
    client.initialize_client(pool_size=100)
    
    # Start QPS statistics thread
    qps_thread = threading.Thread(target=print_qps)
    qps_thread.daemon = True
    qps_thread.start()
    
    # Start multiple sending threads
    threads = []
    for i in range(num_threads):
        thread = threading.Thread(target=send_data, args=(client, target_url))
        thread.daemon = True
        thread.start()
        threads.append(thread)
    
    print(f"Client started with {num_threads} threads, press Ctrl+C to stop")
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("\nClient stopping...")

def main():
    parser = argparse.ArgumentParser(description='Mooncake Communication Bandwidth Test Tool')
    parser.add_argument('mode', choices=['server', 'client'], 
                       help='Run mode: server or client')
    parser.add_argument('--url', default='127.0.0.1:9004',
                       help='URL address (default: 127.0.0.1:9004)')
    parser.add_argument('--threads', type=int, default=8,
                       help='Number of client threads (default: 8)')
    parser.add_argument('--data-size', type=int, default=1,
                       help='Data packet size in MB (default: 1)')

    args = parser.parse_args()
    
    if args.mode == 'server':
        # Server mode, URL as bind address
        bind_url = f"0.0.0.0:{args.url.split(':')[-1]}" if ':' in args.url else f"0.0.0.0:{args.url}"
        run_server(bind_url, args.data_size)
    else:
        # Client mode, URL as target address
        run_client(args.url, args.threads, args.data_size)

if __name__ == "__main__":
    main()
