import json
import zmq
import numpy
import sys
import signal
import base64
import time

sys.path.insert(0, "/opt/dectris/albula/4.1/bin")
sys.path.insert(0, "/opt/dectris/albula/4.1/python")

import dectris.albula as albula
from dectris.albula import DNoObject, DDrawingEllipse, DDrawingString, DDrawingLine

TERMINATE = False


def signal_handler(signal, frame):
    global TERMINATE
    TERMINATE = True


signal.signal(signal.SIGINT, signal_handler)

context = zmq.Context()
socket = context.socket(zmq.SUB)
socket.RCVTIMEO = 1000  # in milliseconds
socket.connect("tcp://xbl-daq-38:5400")
socket.setsockopt(zmq.SUBSCRIBE, b"")

OPTIONAL_DATA = albula.DImageOptionalData()

OPTIONAL_DATA.set_x_pixel_size(0.000075)
OPTIONAL_DATA.set_y_pixel_size(0.000075)

albulaMain = albula.openMainFrame()
albulaSubFrame = albulaMain.openSubFrame()
startTime = time.time()

while not TERMINATE:
    try:
        msg = socket.recv_json()
        if round(time.time() - startTime) >= 60 * 60:
            albulaSubFrame.close()
            albulaSubFrame = albulaMain.openSubFrame()
            startTime = time.time()
        OPTIONAL_DATA.set_wavelength(msg["wavelength_A"])
        OPTIONAL_DATA.set_beam_center_x(msg["beam_x_pxl"])
        OPTIONAL_DATA.set_beam_center_y(msg["beam_y_pxl"])
        OPTIONAL_DATA.set_detector_distance(msg["detector_distance_mm"] * 0.001)
        OPTIONAL_DATA.set_saturation_value(int(msg["saturation_value"]))
        if int(msg["pixel_depth"]) == 2:
            image_array = numpy.frombuffer(base64.b64decode(msg["data"]), numpy.int16)
        else:
            image_array = numpy.frombuffer(base64.b64decode(msg["data"]), numpy.int32)
        height = int(msg["height"])
        width = int(msg["width"])
        image_number = int(msg["image_number"])
        image_array = numpy.reshape(image_array, (height, width))
        print("Received image %d %d %d %d" % (image_number, height, width, len(image_array)))
        dimage = albula.DImage(image_array)
        if image_number >= 0:
            dimage.setOptionalData(OPTIONAL_DATA)
            albulaSubFrame.loadImage(dimage)
            albulaSubFrame.setTitle("JUNGFRAU PREVIEW Dataset: %s Image: %d" % (msg["file_prefix"], image_number))
        else:
            OPTIONAL_DATA.set_pixel_mask(dimage)
    except:
        pass