#!/usr/bin/env python3

import os
import pprint
import sys
import time

from pydicom.dataset import Dataset

from pynetdicom import AE, evt, StoragePresentationContexts, debug_logger
from pynetdicom.sop_class import (
    PatientRootQueryRetrieveInformationModelFind, 
    PatientRootQueryRetrieveInformationModelMove,
    StudyRootQueryRetrieveInformationModelMove,
    PatientStudyOnlyQueryRetrieveInformationModelMove,
)

# 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)
#     print(ds)

    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

# Search in nested Python dict
def gen_dict_extract(key, var):
    if hasattr(var,'items'):
        for k, v in var.items():
            if k == key:
                yield v
            if isinstance(v, dict):
                for result in gen_dict_extract(key, v):
                    yield result
            elif isinstance(v, list):
                for d in v:
                    for result in gen_dict_extract(key, d):
                        yield result

                        
def FindElement(identifier, element):
    e = None
    for v in gen_dict_extract(element, identifier.to_json_dict()):
        e = v['Value'][0]   
    return e  

def QueryDCM(PatientID):

    handlers = [
        (evt.EVT_C_STORE, handle_store),
    ]

#     debug_logger()  
    
    StudyInstanceUID_DICT = {}
    
    ae = AE()
    ae.add_requested_context(PatientRootQueryRetrieveInformationModelFind)
    assoc = ae.associate("192.168.10.56", 104,
        ae_title = 'IQWEBX',
        )
    ds = Dataset()
    ds.QueryRetrieveLevel = 'PATIENT'
    ds.PatientID = PatientID
    responses = assoc.send_c_find(ds, PatientRootQueryRetrieveInformationModelFind)        
#     print(responses)
    for (status, identifier) in responses:
#         print(status, identifier)
        if status:
            print('C-FIND query status: 0x{0:04X}'.format(status.Status))
        else:
            print('Connection timed out, was aborted or received invalid response')
            
        if identifier:
            
#             d = {}
            
#             d['StudyDate'        ] = FindElement(identifier, '00080020')
#             d['ModalitiesInStudy'] = FindElement(identifier, '00080061')
#             d['StudyInstanceUID' ] = FindElement(identifier, '0020000D')
#             print(d)
            
#             if (d['ModalitiesInStudy' ] not in StudyInstanceUID_DICT) or (
#                 d['StudyDate'        ] > StudyInstanceUID_DICT[d['ModalitiesInStudy']]['StudyDate']):
#                 StudyInstanceUID_DICT[d['ModalitiesInStudy']] = d


#         print(dir(identifier))
#         print(identifier.keys())
    


            for sq in identifier[(0x3109, 0x1035)]:
#                 print(sq)
#                 print(type(sq))
    #             exit()
                d = {}

                d['StudyDate'        ] = FindElement(sq, '00080020')
                d['ModalitiesInStudy'] = FindElement(sq, '00080061')
                d['StudyInstanceUID' ] = FindElement(sq, '0020000D')
                print(d)

                if (d['ModalitiesInStudy' ] not in StudyInstanceUID_DICT) or (
                    d['StudyDate'        ] > StudyInstanceUID_DICT[d['ModalitiesInStudy']]['StudyDate']):
                    StudyInstanceUID_DICT[d['ModalitiesInStudy']] = d
        
            
    assoc.release()
    pprint.pprint(StudyInstanceUID_DICT)
#     exit()
    
    
    
    # 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)

    for ModalitiesInStudy, d in StudyInstanceUID_DICT.items():
        StudyInstanceUID = d['StudyInstanceUID']
        
        # Create out identifier (query) dataset
        ds = Dataset()
    #     ds.QueryRetrieveLevel = 'PATIENT'
    #     ds.PatientID = PatientID
        ds.QueryRetrieveLevel = 'STUDY'
        ds.StudyInstanceUID = StudyInstanceUID
        
        
        # 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)

    #         print(responses)
            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()
