import logging
import threading
from collections import deque
from typing import Dict, Set
import xmltodict

def retrieve_input_output_names(openvino_ir_xml_path):
    with open(openvino_ir_xml_path, 'r') as myfile:
        obj = xmltodict.parse(myfile.read())
        input = []
        output = []

        for dic in obj["net"]["layers"]["layer"]:
            if dic.get("@type") == "Parameter":
                input.append(dic["@name"])
            elif dic.get("@type") == "Result":
                output.append(dic["@name"].split("/")[0])
        return input, output



class AsyncPipeline:

    def __init__(self,
                 ie,
                 xml_model_path,
                 plugin_config,
                 device='CPU',
                 max_num_requests=1):

        self.net = ie.read_network(xml_model_path)
        self.logger = logging.getLogger()

        self.logger.info('Loading network to {} plugin...'.format(device))
        self.exec_net = ie.load_network(network=self.net,
                                        device_name=device,
                                        config=plugin_config,
                                        num_requests=max_num_requests)
        if max_num_requests == 0:
            # ExecutableNetwork doesn't allow creation of additional InferRequests. Reload ExecutableNetwork
            # +1 to use it as a buffer of the pipeline
            self.exec_net = ie.load_network(
                network=self.net,
                device_name=device,
                config=plugin_config,
                num_requests=len(self.exec_net.requests) + 1)
        self.empty_requests = deque(self.exec_net.requests)
        self.completed_request_results = {}
        self.callback_exceptions = {}
        self.event = threading.Event()

    def inference_completion_callback(self, status, callback_args):
        try:
            request, id, meta, preprocessing_meta = callback_args
            if status != 0:
                raise RuntimeError(
                    'Infer Request has returned status code {}'.format(status))
            raw_outputs = {
                key: blob.buffer for key, blob in request.output_blobs.items()
            }
            self.completed_request_results[id] = (raw_outputs, meta,
                                                  preprocessing_meta)
            self.empty_requests.append(request)
        except Exception as e:
            self.callback_exceptions.append(e)
        self.event.set()

    def submit_data(self, inputs, id, meta):
        request = self.empty_requests.popleft()
        if len(self.empty_requests) == 0:
            self.event.clear()
        request.set_completion_callback(
            py_callback=self.inference_completion_callback,
            py_data=(request, id, meta, None))
        request.async_infer(inputs=inputs)

    def get_raw_result(self, id):
        if id in self.completed_request_results:
            return self.completed_request_results.pop(id)
        return None

    def get_result(self, id):
        result = self.get_raw_result(id)
        if result:
            raw_result, meta, preprocess_meta = result
            # return self.model.postprocess(raw_result, preprocess_meta), meta
            return raw_result, meta
        return None

    def is_ready(self):
        return len(self.empty_requests) != 0

    def has_completed_request(self):
        return len(self.completed_request_results) != 0

    def await_all(self):
        for request in self.exec_net.requests:
            request.wait()

    def await_any(self):
        if len(self.empty_requests) == 0:
            self.event.wait()


def parse_devices(device_string):
    colon_position = device_string.find(':')
    if colon_position != -1:
        device_type = device_string[:colon_position]
        if device_type == 'HETERO' or device_type == 'MULTI':
            comma_separated_devices = device_string[colon_position + 1:]
            devices = comma_separated_devices.split(',')
            for device in devices:
                parenthesis_position = device.find(':')
                if parenthesis_position != -1:
                    device = device[:parenthesis_position]
            return devices
    return (device_string,)


def parse_value_per_device(devices: Set[str],
                           values_string: str) -> Dict[str, int]:
    """Format: <device1>:<value1>,<device2>:<value2> or just <value>"""
    values_string_upper = values_string.upper()
    result = {}
    device_value_strings = values_string_upper.split(',')
    for device_value_string in device_value_strings:
        device_value_list = device_value_string.split(':')
        if len(device_value_list) == 2:
            if device_value_list[0] in devices:
                result[device_value_list[0]] = int(device_value_list[1])
        elif len(device_value_list) == 1 and device_value_list[0] != '':
            for device in devices:
                result[device] = int(device_value_list[0])
        elif device_value_list[0] != '':
            raise RuntimeError(f'Unknown string format: {values_string}')
    return result

def get_user_config(flags_d: str, flags_nstreams: str,
                    flags_nthreads: int) -> Dict[str, str]:
    config = {}

    devices = set(parse_devices(flags_d))

    device_nstreams = parse_value_per_device(devices, flags_nstreams)
    for device in devices:
        if device == 'CPU':  # CPU supports a few special performance-oriented keys
            # limit threading for CPU portion of inference
            if flags_nthreads:
                config['CPU_THREADS_NUM'] = str(flags_nthreads)

            config['CPU_BIND_THREAD'] = 'NO'

            # for CPU execution, more throughput-oriented execution via streams
            config['CPU_THROUGHPUT_STREAMS'] = str(device_nstreams[device]) \
                if device in device_nstreams else 'CPU_THROUGHPUT_AUTO'
        elif device == 'GPU':
            config['GPU_THROUGHPUT_STREAMS'] = str(device_nstreams[device]) \
                if device in device_nstreams else 'GPU_THROUGHPUT_AUTO'
            if 'MULTI' in flags_d and 'CPU' in devices:
                # multi-device execution with the CPU + GPU performs best with GPU throttling hint,
                # which releases another CPU thread (that is otherwise used by the GPU driver for active polling)
                config['GPU_PLUGIN_THROTTLE'] = '1'
    return config