import glob
import multiprocessing
import os
import subprocess
import sys
import time
import unittest

if sys.platform == "darwin" and multiprocessing.get_start_method(True) is None:
    multiprocessing.set_start_method("fork")

import odil

class Generator(odil.GetSCP.DataSetGenerator):
    def __init__(self):
        odil.GetSCP.DataSetGenerator.__init__(self)
        self._responses = []
        self._response_index = None
        
    def initialize(self, message):
        data_set_1 = odil.DataSet()
        data_set_1.add("SOPClassUID", [odil.registry.RawDataStorage])
        data_set_1.add(
            "SOPInstanceUID", 
            ["1.2.826.0.1.3680043.9.5560.3127449359877365688774406533090568532"])
        data_set_1.add("PatientName", ["Hello^World"])
        data_set_1.add("PatientID", ["1234"])
        self._responses.append(data_set_1)
        
        data_set_2 = odil.DataSet()
        data_set_2.add("SOPClassUID", [odil.registry.RawDataStorage])
        data_set_2.add(
            "SOPInstanceUID", 
            ["1.2.826.0.1.3680043.9.5560.3221615743193123463515381981101110692"])
        data_set_2.add("PatientName", ["Doe^John"])
        data_set_2.add("PatientID", ["5678"])
        self._responses.append(data_set_2)
        
        self._response_index = 0
    
    def done(self):
        return (self._response_index == len(self._responses))
    
    def next(self):
        self._response_index += 1
    
    def get(self):
        return self._responses[self._response_index]
    
    def count(self):
        return 2

class TestGetSCP(unittest.TestCase):
    def test_get_scp_release(self):
        process = multiprocessing.Process(target=self.run_server)
        process.start()
        time.sleep(0.5)
        data_sets = self.run_client()
        process.join(2)
        server_code = process.exitcode
        
        self.assertEqual(server_code, 0)
        
        self.assertEqual(len(data_sets), 2)
        
        self.assertEqual(len(data_sets[0]), 4)
        
        self.assertSequenceEqual(
            data_sets[0].as_string("SOPClassUID"), 
            [odil.registry.RawDataStorage])
        self.assertSequenceEqual(
            data_sets[0].as_string("SOPInstanceUID"), 
            [b"1.2.826.0.1.3680043.9.5560.3127449359877365688774406533090568532"])
        self.assertSequenceEqual(
            data_sets[0].as_string("PatientName"), [b"Hello^World"])
        self.assertSequenceEqual(data_sets[0].as_string("PatientID"), [b"1234"])
        
        self.assertSequenceEqual(
            data_sets[1].as_string("SOPClassUID"), 
            [odil.registry.RawDataStorage])
        self.assertSequenceEqual(
            data_sets[1].as_string("SOPInstanceUID"), 
            [b"1.2.826.0.1.3680043.9.5560.3221615743193123463515381981101110692"])
        self.assertSequenceEqual(
            data_sets[1].as_string("PatientName"), [b"Doe^John"])
        self.assertSequenceEqual(data_sets[1].as_string("PatientID"), [b"5678"])
    
    def run_client(self):
        command = [
            "getscu", 
            "-ll", "error",
            "-P", "-k", "QueryRetrieveLevel=PATIENT",
            "-k", "PatientID=*", "-k", "PatientName",
            "+B",
            "localhost", "11113"]
        
        retcode = subprocess.call(command)
        if retcode != 0:
            return []
        
        files = sorted(glob.glob("{}*".format(odil.uid_prefix)))
        data_sets = []
        for file_ in files:
            with odil.open(file_, "rb") as fd:
                data_sets.append(odil.Reader.read_file(fd)[1])
        for file_ in files:
            os.remove(file_)
        
        return data_sets

    def run_server(self):
        association = odil.Association()
        association.set_tcp_timeout(1)
        association.receive_association("v4", 11113)

        get_scp = odil.GetSCP(association)
        generator = Generator()
        get_scp.set_generator(generator)

        message = association.receive_message()
        get_scp(message)
        
        termination_ok = False

        try:
            association.receive_message()
        except odil.AssociationReleased:
            termination_ok = True
        except odil.AssociationAborted:
            pass
        
        if termination_ok:
            sys.exit(0)
        else:
            sys.exit(1)

if __name__ == "__main__":
    unittest.main()
