#!/usr/bin/env python3

import os
import sys
import time

from pydicom.dataset import Dataset

from pynetdicom import AE, evt, StoragePresentationContexts, debug_logger
from pynetdicom.sop_class import PatientRootQueryRetrieveInformationModelMove

# debug_logger()

Series = {}

dcm_dir = '/nn'

def handle_store(event):
    # print(event)
    """Handle a C-STORE request event."""
    ds = event.dataset
    ds.file_meta = event.file_meta

    # print(ds.SeriesInstanceUID)

    if ds.SeriesInstanceUID in Series:
        s = Series[ds.SeriesInstanceUID]
    else:
        number = len(Series)         
        series_dir = os.path.join(dcm_dir, ds.PatientID, ds.StudyInstanceUID, ds.SeriesInstanceUID, str(number))
        series_dir = os.path.join(dcm_dir, ds.PatientID, ds.StudyDate, ds.Modality, str(number))
        series_dir = os.path.join(dcm_dir, ds.PatientID, ds.StudyDate, ds.Modality, hex(number)[2:])
        s = {
            'PatientID': ds.PatientID,
            'StudyInstanceUID': ds.StudyInstanceUID,
            'SeriesInstanceUID': ds.SeriesInstanceUID,

            'Modality': ds.Modality,
            'SeriesDescription': ds.SeriesDescription,

            'path': series_dir,
            'number': number,
        }
        # print(s)
        print(s['path'], s['SeriesDescription'])
        Series[ds.SeriesInstanceUID] = s

    os.makedirs(s['path'], exist_ok=True)

    filename = os.path.join(s['path'], ds.SOPInstanceUID)

    # Save the dataset using the SOP Instance UID as the filename
    ds.save_as(filename, write_like_original=False)

    # Return a 'Success' status
    return 0x0000

def QueryDCM(PatientID):

    handlers = [(evt.EVT_C_STORE, handle_store)]

    # Initialise the Application Entity
    ae = AE()

    # Add a requested presentation context
    ae.add_requested_context(PatientRootQueryRetrieveInformationModelMove)

    # Add the Storage SCP's supported presentation contexts
    ae.supported_contexts = StoragePresentationContexts

    # Start our Storage SCP in non-blocking mode, listening on port 11120
    ae.ae_title = 'OUR_STORE_SCP'
    # scp = ae.start_server(("127.0.0.1", 11120), block=False, evt_handlers=handlers)
    scp = ae.start_server(("0.0.0.0", 11120), block=False, evt_handlers=handlers)

    # Create out identifier (query) dataset
    ds = Dataset()
    ds.QueryRetrieveLevel = 'PATIENT'
    ds.PatientID = PatientID


    # Associate with peer AE at IP 127.0.0.1 and port 11112
    # assoc = ae.associate("127.0.0.1", 11112)
    assoc = ae.associate("192.168.10.56", 104,
        ae_title = 'IQWEBX',
        )


    if assoc.is_established:
        # Use the C-MOVE service to send the identifier
        responses = assoc.send_c_move(ds, 'OUR_STORE_SCP', PatientRootQueryRetrieveInformationModelMove)

        for (status, identifier) in responses:
            # print(status, identifier)
            if status:
                print('C-MOVE query status: 0x{0:04x}'.format(status.Status))
            else:
                print('Connection timed out, was aborted or received invalid response')

        # Release the association
        assoc.release()
    else:
        print('Association rejected, aborted or never connected')

    # Stop our Storage SCP
    scp.shutdown()


def main():
    if len(sys.argv) < 2:
        print('Usage:', sys.argv[0], 'PatientID')
        sys.exit()
    print('hello')
    print(sys.argv[0])
    print(sys.argv[1])

    start = time.time()
    QueryDCM(sys.argv[1])
    end = time.time()
    print(end - start, 'seconds')

    # for k, v in Series.items():
        # print(v['number'], v['Modality'], v['SeriesDescription'], v['path'])
        # print(v['path'], v['SeriesDescription'])

if __name__ == '__main__':
    main()