import pymoku

log = pymoku.log.getChild('sources')


class RoutingException(Exception):
    def __init__(self, message, conflicts):
        super().__init__(message)
        self.conflicts = conflicts


class Edge(object):
    def __init__(self, node=None, idx=None, bus_width=1, limit=None):
        if not isinstance(idx, tuple):
            idx = (idx,)
        self.node = node
        self.idx = idx
        self.bus_width = bus_width
        self.limit = limit

    def attach(self, node, owner=None, bus_width=1):
        if self.bus_width != bus_width:
            return False

        # limit restricts this edge to only the listed upstream leaves
        # TODO this limits a specific target node, not upstream branches
        if self.limit is not None:
            for lnode in self.limit:
                if node is lnode or node == lnode.name:
                    break
            else:
                return False

        return self.node.attach(node, owner=owner, bus_width=bus_width)

    def __repr__(self):
        return f"{type(self).__name__}(node='{self.node.name}', idx={self.idx})"


class Node(object):
    def __init__(self, name='', set_idx=None, get_idx=None):
        self.name = name

        if set_idx is not None:
            self.set_idx = set_idx

        if get_idx is not None:
            self.get_idx = get_idx

        self.clear()

    def clear(self):
        self._input_edges = []
        self._refs = set()
        self.idx = None
        self._conflicts = set()

    def add(self, node=None, name=None, **kwargs):
        if node is None:
            node = Node(name)

        # idx to node is not required to be one-to-one
        self._input_edges.append(Edge(node, **kwargs))
        return node  # convienience

    def attach(self, node, owner=None, bus_width=1):
        if owner is None:
            self._conflicts = set()

        if (
            node is self
            or (self.name == node)
        ):
            self._refs.add(owner)
            return True

        self.release(owner)

        if node is None:
            self.set_idx(None)
            return True

        for edge in self._input_edges:
            if edge.attach(node, owner=owner or self, bus_width=bus_width):
                if self.get_idx() == edge.idx:
                    # already selected and shared
                    log.debug(f'attach: {edge} [{bus_width:d}]-> {self}')
                    self._refs.add(owner)
                    return True

                if not self._refs:
                    self._refs.add(owner)
                    log.debug(f'select: {edge} [{bus_width:d}]-> {self}')
                    self.set_idx(edge.idx)
                    return True

                owner.add_conflicts(self._refs)
                edge.node.release(self)
                continue

        return False

    def active_edges(self):
        idx = self.get_idx()
        return [edge for edge in self._input_edges if edge.idx == idx]

    def get_attached(self, limit=None):
        if not self._input_edges:
            return self

        for edge in self.active_edges():
            # we accumulate all limits along the path into one set
            # then filter this edge with only the downstream limits
            new_limit = set()
            new_limit.update(limit or {})
            new_limit.update(edge.limit or {})
            node = edge.node.get_attached(new_limit)
            if node is not None:
                if not limit or node in limit:
                    return node
        return None

    def set_idx(self, idx):
        self.idx = idx

    def get_idx(self):
        return self.idx

    def release(self, owner=None):
        for edge in self._input_edges:
            edge.node.release(owner or self)
        if owner in self._refs:
            self._refs.remove(owner)

    def find_node(self, name):
        if self.name == name:
            return self

        for edge in self._input_edges:
            n = edge.node.find_node(name)
            if n is not None:
                return n

        return None

    def paths_to(self, node, bus_width=None):
        paths = []
        for edge in self._input_edges:
            if bus_width is not None and edge.bus_width != bus_width:
                continue

            if edge.limit is not None and node not in edge.limit:
                continue

            if edge.node is node:
                paths.append([(self, edge)])
                continue

            up_paths = edge.node.paths_to(node, bus_width)
            for up_path in up_paths:
                paths.append(up_path + [(self, edge)])

        return paths

    def add_conflicts(self, conflict):
        for conf in conflict:
            self._conflicts.add(conf)

    def __repr__(self):
        return f"{type(self).__name__}(name='{self.name}')"

    def graph(self, dot=None, annotations=False):
        if dot is None:
            import graphviz  # noqa
            dot = graphviz.Digraph(engine='dot', strict=annotations)
            dot.graph_attr['rankdir'] = 'LR'
            dot.format = 'svg'

        # skip if this node already exists in the graph
        if f"\t{self.name}" in dot.source:
            return

        for edge in self._input_edges:
            edge.node.graph(dot=dot, annotations=annotations)

        if not self._input_edges:
            args = dict(style='filled', fillcolor='#40e0d0')
        else:
            args = dict()

        dot.node(self.name, **args)

        for edge in self._input_edges:
            if annotations:
                dot.edge(edge.node.name,
                         self.name,
                         label=f'{edge.idx}' + (f' {list(l.name for l in edge.limit)}' if edge.limit else ''),
                         penwidth=f'{edge.bus_width:d}',
                         color='red' if edge.node.get_idx() == edge.idx else 'black'
                         )
            else:
                dot.edge(edge.node.name, self.name)
        return dot


class Input(Node):
    def __init__(self, name, slot_inst, bus_width=1, set_idx=None, get_idx=None):
        super().__init__(name=name, set_idx=set_idx, get_idx=get_idx)
        self.slot_inst = slot_inst
        self.moku = slot_inst.moku
        self.bus_width = bus_width
        self._callbacks = []

    def get_source(self):
        with self.moku.cached_state(), self.slot_inst.cached_registers():
            return self.get_attached()

    def set_source(self, source):
        if not isinstance(source, Node):
            source = self.find_node(source)

        with self.moku.cached_state(), self.slot_inst.cached_registers():
            bus_width = min(self.bus_width, source.bus_width) if source else self.bus_width
            if not self.attach(source, bus_width=bus_width):
                src_name = source if isinstance(source, str) else source.name
                raise RoutingException(f"Couldn't resolve route {src_name} -> {self.name}",
                                       conflicts=self._conflicts)

            for cbk in self._callbacks:
                cbk(self, source)

    def coeff(self):
        if (node := self.get_attached()) is None or node is self:
            return 1.0

        try:
            return node.coeff()
        except TypeError:  # not callable
            return node.coeff
        except AttributeError:  # no coeff
            raise Exception(f"Node {node} has no coeff")

    def slot(self):
        return self.slot_inst

    def on_change(self, callback):
        self._callbacks.append(callback)


class Output(Node):
    def __init__(self, name, slot_inst, bus_width=1, coeff=2**16, dtype=pymoku.dtype.s16x4):
        super().__init__(name=name)
        self.slot_inst = slot_inst
        self.bus_width = bus_width
        self._coeff = coeff
        self.dtype = dtype

    def coeff(self):
        return self._coeff

    def slot(self):
        return self.slot_inst


class ADC(Node):
    def __init__(self, name, coeff, bus_width=1, dtype=pymoku.dtype.s16x4):
        super().__init__(name=name)
        self.coeff = coeff
        self.bus_width = bus_width
        self.dtype = dtype


class DAC(Node):
    def __init__(self, name, routing, dac_idx, coeff=2**16, bus_width=1):
        super().__init__(name=name)
        self.routing = routing
        self.dac_idx = dac_idx
        self.coeff = coeff
        self.bus_width = bus_width

    def set_idx(self, idx):
        if idx is None:
            idx = [0xFF] * self.bus_width
        self.routing.dacs[self.dac_idx:self.dac_idx + self.bus_width] = idx

    def get_idx(self):
        return tuple(self.routing.dacs[self.dac_idx:self.dac_idx + self.bus_width])

    def get_source(self):
        return self.get_attached()

    def set_source(self, source):
        bus_width = min(source.bus_width, self.bus_width) if source is not None else 1
        if not self.attach(source, bus_width=bus_width):
            raise RoutingException(f"Couldn't resolve route {source.name} -> {self.name}",
                                   conflicts=self._conflicts)


class LoopMux(Node):
    def __init__(self, slot_ctrl, loop_idx, width=1):
        super().__init__(name=f'LoopMux{width:d}X_{slot_ctrl.slot_id:d}_{loop_idx:d}')
        self.slot_ctrl = slot_ctrl
        self.loop_idx = loop_idx
        self.width = width

    def get_idx(self):
        return tuple(self.slot_ctrl.loops[self.loop_idx:self.loop_idx + self.width])

    def set_idx(self, idx):
        if idx is None:
            idx = [0xFF] * self.width
        self.slot_ctrl.loops[self.loop_idx:self.loop_idx + self.width] = idx


class SlotOutput(Node):
    def __init__(self, slot, outp_idx):
        super().__init__(name=f'SlotOutput{slot:d}_{outp_idx:d}')
        self.slot_idx = slot
        self.outp_idx = outp_idx

    def slot(self):
        # if we resolved here, there's probably nothing deployed
        return None

    def get_attached(self, limit=None):
        if self._input_edges:
            return self._input_edges[0].node.get_attached(limit)
        return self
