from debug import trace as grownet_trace


class Tract:
    """Bridges two layers by subscribing to source fires and routing to dest.

    Extended to optionally pass 2D context and/or deterministic sinks.
    """
    def __init__(self, src_layer, dst_layer, region_bus=None, feedback=False,
                 probability: float | None = None,
                 allowed_source_indices: set[int] | None = None,
                 sink_map: dict[int, list[int] | set[int]] | None = None):
        self.src = src_layer
        self.dst = dst_layer
        self.region_bus = region_bus
        self.feedback = bool(feedback)
        self.sink_map = sink_map or {}
        self.allowed = allowed_source_indices  # if None: allow all

        # capture source shape if 2D
        self.src_height = getattr(self.src, "height", None)
        self.src_width = getattr(self.src, "width", None)

        # subscribe to source neuron fires
        for src_index, neuron in enumerate(self.src.get_neurons()):
            if self.allowed is not None and src_index not in self.allowed:
                continue
            def make_hook(i):
                return lambda who, value: self.on_source_fired(i, value)
            neuron.register_fire_hook(make_hook(src_index))

    def on_source_fired(self, source_index, value):
        # If we have an explicit sink map (e.g., windowed wiring to OutputLayer2D), deliver directly
        targets = self.sink_map.get(source_index)
        if targets:
            try:
                neurons = self.dst.get_neurons()
                # Deduplicate in case a center was added multiple times.
                for target_index in (targets if isinstance(targets, set) else set(targets)):
                    if 0 <= target_index < len(neurons):
                        target_neuron = neurons[target_index]
                        if hasattr(self.dst, "region") and self.dst.region is not None:
                            self.dst.region._record_delivery_event()
                        delivery_event = {}
                        delivery_event["source_neuron_id"] = self.source_neuron_id(source_index)
                        delivery_event["target_neuron_id"] = getattr(target_neuron, "id", None)
                        delivery_event["output_value"] = float(value)
                        delivery_event["reason"] = "windowed"
                        grownet_trace.emit("tract.deliver", **delivery_event)
                        with grownet_trace.run_with_hop_increment():
                            fired = target_neuron.on_input(value)
                        if fired:
                            target_neuron.on_output(value)
                return
            except Exception:
                pass

        # Otherwise, prefer 2D-aware propagation if shape is known and dst supports it
        if self.src_height is not None and self.src_width is not None and hasattr(self.dst, "propagate_from_2d"):
            delivery_event = {}
            delivery_event["source_neuron_id"] = self.source_neuron_id(source_index)
            delivery_event["output_value"] = float(value)
            delivery_event["reason"] = "windowed_fanout"
            grownet_trace.emit("tract.deliver", **delivery_event)
            # The destination layer records actual per-target deliveries during propagate_from_2d.
            with grownet_trace.run_with_hop_increment():
                self.dst.propagate_from_2d(source_index, value, int(self.src_height), int(self.src_width))
            return

        # Fallback: let the destination layer handle generic propagate
        delivery_event = {}
        delivery_event["source_neuron_id"] = self.source_neuron_id(source_index)
        delivery_event["output_value"] = float(value)
        delivery_event["reason"] = "layer_propagate"
        grownet_trace.emit("tract.deliver", **delivery_event)
        # The destination layer records actual per-target deliveries during propagate_from.
        with grownet_trace.run_with_hop_increment():
            self.dst.propagate_from(source_index, value)

    # ---- growth hook: attach a newly created source neuron ----
    def attach_source_neuron(self, new_src_index: int) -> None:
        if self.allowed is not None and new_src_index not in self.allowed:
            return
        neurons = self.src.get_neurons()
        if not (0 <= new_src_index < len(neurons)):
            return
        def hook(neuron_obj, amplitude):
            self.on_source_fired(new_src_index, amplitude)
        try:
            neurons[new_src_index].register_fire_hook(hook)
        except Exception:
            pass

    def source_neuron_id(self, source_index: int):
        try:
            source_neurons = self.src.get_neurons()
            if 0 <= int(source_index) < len(source_neurons):
                return getattr(source_neurons[int(source_index)], "id", None)
        except Exception:
            pass
        return None
