"""
Record messages
"""

from dataclasses import dataclass
from typing import Any
import tqdm
import zmq
import random
import sys
import time
import json
import jsonlines

port = "99174"
if len(sys.argv) > 1:
    port =  sys.argv[1]
    int(port)

if len(sys.argv) > 2:
    fn = sys.argv[2]
else:
    fn = "messages.jsonl"

context = zmq.Context()
socket = context.socket(zmq.SUB)

addr = f"tcp://100.109.175.82:{port}"

socket.connect (addr)
socket.setsockopt_string(zmq.SUBSCRIBE, "")
#print('connected')
# Subscribe to zipcode, default is NYC, 10001
#topicfilter = "10001"
#socket.setsockopt(zmq.SUBSCRIBE, topicfilter)

last = time.time()

print(f"Listen to {addr}, write to {fn}")

with jsonlines.open(fn, mode='w', flush=True) as writer:
    with tqdm.tqdm() as pbar:
        while True:
            pbar.update()
            string = socket.recv_string()
            data = json.loads(string)        
            now = time.time()
            msg = {
                'data': data,
                'offset': now - last
            }
            writer.write(msg)
            last = now
        
    #topic, messagedata = string.split()
    #total_value += int(messagedata)
    #print topic, messagedata

#print "Average messagedata value for topic '%s' was %dF" % (topicfilter, total_value / update_nbr)