#!/usr/bin/env python3

import sys
import pika
import hashlib

from scipy.stats import chisquare

def get_queue_name(i):
  return str(i)

def load_queues():
  for i in range(num_trials):
    channel.basic_publish(exchange = exchange_name, routing_key = hashlib.sha256(str(i).encode('utf-8')).hexdigest(), body = str(i))

def analyze_queues():
  num_messages = []
  for i in range(num_queues):
    redeclare_result = channel.queue_declare(queue = get_queue_name(i), passive = True)
    num_messages.append(int(redeclare_result.method.message_count))
  print('Message count by queue:')
  print(num_messages)

  print('Chi-Squared test result:')
  print(chisquare(num_messages))

def reset_queues():
  for i in range(num_queues):
    channel.queue_delete(queue = get_queue_name(i))
    channel.exchange_delete(exchange = exchange_name)


username = sys.argv[1]
password = sys.argv[2]
hostname = sys.argv[3]
port = sys.argv[4]
vhost = sys.argv[5]
mode = sys.argv[6]

credentials = pika.PlainCredentials(username, password)
connection = pika.BlockingConnection(pika.ConnectionParameters(hostname, int(port), vhost, credentials))
channel = connection.channel()

exchange_name = 'test_consistent_hash'
num_queues = 10
num_trials = 10000
channel.exchange_declare(exchange = exchange_name, exchange_type = 'x-consistent-hash')

for i in range(num_queues):
  channel.queue_declare(queue = get_queue_name(i))
  channel.queue_bind(exchange = exchange_name, queue = get_queue_name(i), routing_key='1')

if mode == 'load':
  load_queues()
elif mode == 'analyze':
  analyze_queues()
elif mode == 'reset':
  reset_queues()