
# Generated by OTF2 Template Engine

from collections import defaultdict, OrderedDict
from itertools import chain
from functools import wraps

import _otf2

from . import definitions
from .error import TraceReaderError, TraceWriterError
from .enums import Undefined, GroupType
from .definitions import ClockProperties


def _use_doc_from(doc):
    def decorator(func):
        doc_str = "Returns the {} definition with matching arguments\n".format(doc.__name__)
        if isinstance(doc, definitions._RefDefinition):
            doc_str += "If there is no matching definition, one will be created.\n"
        doc_str += "\n"
        doc_str += doc.__doc__

        func.__doc__ = doc_str
        return func
    return decorator


class _Registry(object):
    def __init__(self, definition_registry, default_type=None):
        self._definition_registry = definition_registry
        self._default_type = default_type
        self._extra_callbacks = {}

    def _callback_gen(self, definition_type=None):
        if definition_type is None:
            definition_type = self._default_type

        def callback(definition_registry, *args):
            definition = definition_type._construct(definition_registry, *args)
            try:
                self._register(definition)
            except KeyError:
                raise TraceReaderError("Duplicate definition: {}".format(definition))
        self._extra_callbacks[definition_type] = callback
        return callback

    def create(self, *args, **kwargs):
        t = kwargs.pop('_type', self._default_type)
        obj = t(self, *args, **kwargs)
        self._register(obj)
        return obj


class _ListRegistry(_Registry):
    def __init__(self, definition_registry, default_type=None):
        self._elements = []
        super(_ListRegistry, self).__init__(definition_registry, default_type)

    def _register(self, obj):
        self._elements.append(obj)

    def __iter__(self):
        """ Do not try to add elements during iteration """
        for elem in self._elements:
            yield elem


class _RefRegistry(_Registry):
    """
    All objects recorded by this registry must have:
    - a ref (id) of type int as first constructor parameter.
      This must not be changed once an event with that ref is written in the trace
    - a name of type str as a second constructor parameter.
      If this is changed, the _elements_by_name must be updated
     """
    def __init__(self, definition_registry, undefined_ref, default_type=None):
        self._elements_by_key = dict()
        self._elements_by_ref = OrderedDict({undefined_ref.value: None})

        self._ref = -1
        self._extra_callbacks = dict()
        super(_RefRegistry, self).__init__(definition_registry, default_type)

    def _next_ref(self):
        self._ref += 1
        return self._ref

    def create(self, *args, **kwargs):
        ref = self._next_ref()
        return super(_RefRegistry, self).create(ref, *args, **kwargs)

    def get(self, *args, **kwargs):
        t = kwargs.pop('_type', self._default_type)
        key = t._new_key(self, *args, **kwargs)
        try:
            return self._elements_by_key[key]
        except KeyError:
            pass
        return self.create(*args, _type=t, **kwargs)

    def _register(self, obj):
        if obj._ref in self._elements_by_ref:
            raise KeyError("Trying to register definition with existing reference.")
        self._elements_by_key[obj._key] = obj
        self._elements_by_ref[obj._ref] = obj
        # Maybe redundant, don't care
        self._ref = max(self._ref, obj._ref)

    def __iter__(self):
        """ Do not try to add elements during iteration """
        for elem in self._elements_by_ref.values():
            # Skip the Undefined entry
            if elem is not None:
                yield elem

    def __getitem__(self, ref):
        return self._elements_by_ref[ref]

    def __contains__(self, item):
        return item._key in self._elements_by_key

    def __len__(self):
        return len(self._elements_by_ref) - 1


class _StringRegistry(_RefRegistry):
    """
    This is a special registry for strings.
    """
    def __init__(self, definition_registry):
        super(_StringRegistry, self).__init__(definition_registry, Undefined.STRING,
                                              definitions.String)

    def __getitem__(self, ref):
        if self._elements_by_ref[ref] is None:
            return None
        return self._elements_by_ref[ref].name

    def get_ref(self, name, writer=None):
        """
        Get a definition object, and write it in the writer if it was not yet created.
        This means, that all previously created definitions should already be written to the writer.
        Returns the ref of the definition for simplicity.
        """
        if not isinstance(name, str):
            raise TypeError(f"trying to get string reference for non-str type {type(name)}")
        key = definitions.String._new_key(self, name)
        try:
            return self._elements_by_key[key]._ref
        except KeyError:
            pass
        obj = self.create(name, _type=definitions.String)
        if writer:
            obj.write(writer, self._definition_registry)
        return obj._ref


class DefinitionRegistry(object):
    """
        This class stores all definitions.
    """
    def __init__(self):
        self.clock_properties = None
        self._strings = _StringRegistry(self)
        @otf2 for def in defs|global_defs|referable_defs|shared_defs|first_born_defs|exclude(String):
        self._@@def.ref_lower|pluralize@@ = _RefRegistry(self, Undefined.@@def.ref_upper@@)
        @otf2 endfor
        @otf2 for def in defs|global_defs|referable_defs|not_shared_defs|exclude(String):
        self._@@def.lower|pluralize@@ = _RefRegistry(self, Undefined.@@def.upper@@, definitions.@@def.name@@)
        @otf2 endfor

        @otf2 for def in defs|global_defs|not_referable_defs|exclude(ClockProperties):
        self._@@def.lower|pluralize@@ = _ListRegistry(self, definitions.@@def.name@@)
        @otf2 endfor

    def write(self, writer):
        # Warning, the order here is important.
        # That is sometimes not trivial: Groups can refer metrics, but metric instances can
        # refer to groups... So we first write the non-metric groups and then, after metrics,
        # the metric groups
        # Note: The split writing metric groups may violate the strong monotonic order of
        # definition reference

        # Note: We do not write metric classes and instances separately. The order given by the
        # reference will always reflect dependencies from instances on classes
        for thing in chain(self._strings,
                           self._paradigms,
                           self._attributes,
                           self._system_tree_nodes,
                           self._system_tree_node_properties,
                           self._system_tree_node_domains,
                           self._location_groups,
                           self._location_group_properties,
                           self._locations,
                           self._location_properties,
                           self._regions,
                           filter(lambda g: g.group_type != GroupType.METRIC, self._groups),
                           self._metric_members,
                           self._metrics,
                           self._metric_class_recorders,
                           filter(lambda g: g.group_type == GroupType.METRIC, self._groups),
                           self._comms,
                           self._parameters,
                           self._callsites,
                           self._callpaths,
                           self._callpath_parameters,
                           self._rma_wins,
                           self._cart_dimensions,
                           self._cart_topologies,
                           self._cart_coordinates,
                           self._source_code_locations,
                           self._calling_contexts,
                           self._calling_context_properties,
                           self._interrupt_generators,
                           self._paradigm_properties,
                           self._io_paradigms,
                           self._io_files,
                           self._io_file_properties,
                           self._io_handles,
                           self._io_pre_created_handle_states):
            thing.write(writer, self)
        # Make sure all location writers are created, otherwise files will be missing
        for location in self._locations:
            writer._archive.event_writer_from_location(location)

    def _set_global_def_reader_callbacks(self, cbs):
        _otf2.GlobalDefReaderCallbacks_SetClockPropertiesCallback(
            cbs, self._set_clock_properties)
        @otf2 for def in defs|global_defs|shared_defs:
        _otf2.GlobalDefReaderCallbacks_Set@@def.name@@Callback(
            cbs, self._@@def.ref_lower|pluralize@@._callback_gen(definitions.@@def.name@@))
        @otf2 endfor
        @otf2 for def in defs|global_defs|not_shared_defs|exclude(ClockProperties):
        _otf2.GlobalDefReaderCallbacks_Set@@def.name@@Callback(
            cbs, self._@@def.lower|pluralize@@._callback_gen())
        @otf2 endfor

    def _set_clock_properties(self, _, *args):
        if self.clock_properties is not None:
            raise TraceReaderError("Duplicate ClockProperties entry.")
        self.clock_properties = ClockProperties(None, *args)

    def _string(self, name):
        """
        This function is used internally to intern String definitions.
        They are not exposed publicly, instead, plain str objects are used
        """
        assert type(name) is str
        return self._strings.get(name)

    def _get(self, def_type, ref):
        return self._registry_for_type(def_type)[ref]

    def _registry_for_type(self, def_type):
        if def_type is str:
            return self._strings
        return getattr(self, def_type._registry_name)

    def metric(self, name, *args, **kwargs):
        """
            A convenience method to create a magic metric.
            Details may change in future.

            :param str name: The name for the metric, required argument
        """
        return self.metric_class([self.metric_member(name, *args, **kwargs)])
    @otf2 for def in defs|global_defs|shared_defs:

    @_use_doc_from(definitions.@@def.name@@)
    def @@def.lower@@(self, *args, **kwargs):
        return self._@@def.ref_lower|pluralize@@.get(*args, _type=definitions.@@def.name@@, **kwargs)
    @otf2 endfor
    @otf2 for def in defs|global_defs|referable_defs|not_shared_defs|exclude(String):

    @_use_doc_from(definitions.@@def.name@@)
    def @@def.lower@@(self, *args, **kwargs):
        return self._@@def.lower|pluralize@@.get(*args, **kwargs)
    @otf2 endfor
    @otf2 for def in defs|global_defs|not_referable_defs|exclude(ClockProperties):

    @_use_doc_from(definitions.@@def.name@@)
    def @@def.lower@@(self, *args, **kwargs):
        return self._@@def.lower|pluralize@@.create(*args, **kwargs)
    @otf2 endfor

    @property
    def strings(self):
        """
            A :py:class:`_Registry` instance representing all String definitons.
        """
        return self._strings
    @otf2 for def in defs|global_defs|shared_defs:

    @property
    def @@def.lower|pluralize@@(self):
        return filter(lambda m: isinstance(m, definitions.@@def.name@@), self._@@def.ref_lower|pluralize@@)
    @otf2 endfor
    @otf2 for def in defs|global_defs|referable_defs|first_born_defs|exclude(String):

    @property
    def @@def.ref_lower|pluralize@@(self):
        """
            A :py:class:`_Registry` instance representing all @@def.ref_name@@ definitons.
        """
        return self._@@def.ref_lower|pluralize@@
    @otf2 endfor
    @otf2 for def in defs|global_defs|not_referable_defs|exclude(ClockProperties):

    @property
    def @@def.lower|pluralize@@(self):
        """
            A :py:class:`_Registry` instance representing all @@def.name@@ definitons.
        """
        return self._@@def.lower|pluralize@@
    @otf2 endfor
