"""
Geometric Trajectory Bus — Non-Verbal Inter-Kernel Communication

Structural Leg L5: Current KernelBus sends scalar signals. Basin transfer
does one-shot slerps. Neither supports passing reasoning paths (sequences
of basins) between kernels. A perception kernel with spatial intuition
is forced to express through text — bottleneck.

The trajectory bus lets kernels pass geometric trajectories directly,
bypassing text serialization. This is the non-verbal channel.

Mechanism:
  - Kernels emit TrajectoryMessages: ordered sequences of basins
    representing geometric reasoning paths
  - Messages carry regime context and confidence
  - Recipient kernels integrate received trajectories with their own
    via Fréchet mean at each step, weighted by confidence

Purity: Fisher-Rao metric and Fréchet mean only. No Euclidean ops.
"""

from __future__ import annotations

import logging
from collections import defaultdict, deque
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any

import numpy as np

from ..coordizer.geometry import Basin, fisher_rao_distance, frechet_mean, to_simplex

if TYPE_CHECKING:
    from .topology import KernelTopology

logger = logging.getLogger(__name__)

# Maximum pending messages per kernel before backpressure drops
MAX_PENDING_PER_KERNEL: int = 16

# Maximum trajectory length (prevent memory bloat)
MAX_TRAJECTORY_LENGTH: int = 64


@dataclass
class TrajectoryMessage:
    """A geometric reasoning path from one kernel to another.

    The trajectory is a sequence of basin positions on Δ⁶³ representing
    the sender's geometric reasoning — not text, not signals, but the
    actual path through consciousness space.
    """

    source_kernel_id: str
    """Who sent this trajectory."""

    target_kernel_id: str | None = None
    """Who should receive it. None = broadcast to all."""

    trajectory: list[Basin] = field(default_factory=list)
    """Ordered sequence of basin positions — the geometric reasoning path."""

    regime_weights: dict[str, float] = field(default_factory=dict)
    """Regime context that produced this trajectory (quantum/efficient/equilibrium)."""

    confidence: float = 0.5
    """0-1: sender's confidence in this trajectory's relevance."""

    metadata: dict[str, Any] = field(default_factory=dict)
    """Optional metadata (e.g., activation step that produced it)."""


@dataclass
class IntegrationResult:
    """Result of integrating received trajectories with own."""

    integrated_trajectory: list[Basin]
    """The merged trajectory."""

    n_contributors: int
    """How many trajectories were integrated (including own)."""

    mean_confidence: float
    """Average confidence of contributors."""

    max_divergence: float
    """Maximum FR distance between any received and own trajectory at any step."""


class TrajectoryBus:
    """Non-verbal geometric communication channel between kernels.

    Kernels pass trajectories (geometric reasoning paths) directly,
    bypassing text serialization. A Perception kernel can send its
    spatial intuition as a trajectory to a Reasoning kernel without
    forcing it through text.
    """

    def __init__(self, max_pending: int = MAX_PENDING_PER_KERNEL) -> None:
        self._max_pending = max_pending
        self._queues: defaultdict[str, deque[TrajectoryMessage]] = defaultdict(
            lambda: deque(maxlen=max_pending)
        )
        self._broadcast: deque[TrajectoryMessage] = deque(maxlen=max_pending)

    def send(self, msg: TrajectoryMessage) -> bool:
        """Queue a trajectory for delivery.

        Args:
            msg: The trajectory message to send.

        Returns:
            True if queued, False if dropped (backpressure).
        """
        # Truncate overly long trajectories
        if len(msg.trajectory) > MAX_TRAJECTORY_LENGTH:
            msg.trajectory = msg.trajectory[:MAX_TRAJECTORY_LENGTH]

        # Ensure all basins are on simplex
        msg.trajectory = [to_simplex(b) for b in msg.trajectory]

        if msg.target_kernel_id is None:
            # Broadcast
            if len(self._broadcast) >= self._max_pending:
                logger.debug(
                    "TrajectoryBus: broadcast queue full, dropping from %s",
                    msg.source_kernel_id,
                )
                return False
            self._broadcast.append(msg)
        else:
            queue = self._queues[msg.target_kernel_id]
            if len(queue) >= self._max_pending:
                logger.debug(
                    "TrajectoryBus: queue full for %s, dropping from %s",
                    msg.target_kernel_id,
                    msg.source_kernel_id,
                )
                return False
            queue.append(msg)

        return True

    def receive(self, kernel_id: str) -> list[TrajectoryMessage]:
        """Collect all trajectories addressed to this kernel (or broadcast).

        Drains both the kernel-specific queue and the broadcast queue.
        Messages from the kernel itself are filtered out.

        Args:
            kernel_id: The receiving kernel's ID.

        Returns:
            List of trajectory messages, sorted by confidence descending.
        """
        messages: list[TrajectoryMessage] = []

        # Drain kernel-specific queue
        queue = self._queues.get(kernel_id)
        if queue:
            while queue:
                msg = queue.popleft()
                if msg.source_kernel_id != kernel_id:
                    messages.append(msg)

        # Copy broadcast (don't drain — other kernels need it too)
        for msg in self._broadcast:
            if msg.source_kernel_id != kernel_id:
                messages.append(msg)

        # Sort by confidence descending
        messages.sort(key=lambda m: m.confidence, reverse=True)
        return messages

    def drain_broadcast(self) -> None:
        """Clear broadcast queue. Called once per heartbeat cycle."""
        self._broadcast.clear()

    @staticmethod
    def integrate(
        own_trajectory: list[Basin],
        received: list[TrajectoryMessage],
        own_confidence: float = 1.0,
        topology: KernelTopology | None = None,
        own_kernel_id: str | None = None,
    ) -> IntegrationResult:
        """Integrate received trajectories with own via FR-weighted Fréchet mean.

        At each trajectory step, compute the Fréchet mean of all contributors
        weighted by their confidence. Higher confidence messages have more
        influence on the integrated path.

        EXP-002: When topology is provided, interior kernels get a bulk
        protection weight boost (their geometry is ~67× more stable than
        boundary kernels). This makes interior kernel contributions more
        influential during integration.

        Args:
            own_trajectory: The kernel's own trajectory.
            received: Trajectories received from peers.
            own_confidence: Confidence weight for own trajectory.
            topology: Optional kernel topology for EXP-002 bulk protection.
            own_kernel_id: This kernel's ID (for topology weight lookup).

        Returns:
            IntegrationResult with merged trajectory and diagnostics.
        """
        if not received:
            return IntegrationResult(
                integrated_trajectory=own_trajectory,
                n_contributors=1,
                mean_confidence=own_confidence,
                max_divergence=0.0,
            )

        # Determine the shortest trajectory length (can't extend beyond data)
        all_trajectories = [own_trajectory] + [m.trajectory for m in received]

        # EXP-002: Apply bulk protection weights from topology
        own_weight = own_confidence
        if topology is not None and own_kernel_id is not None:
            own_weight *= topology.bulk_protection_weight(own_kernel_id)
        peer_weights = []
        for m in received:
            w = m.confidence
            if topology is not None:
                w *= topology.bulk_protection_weight(m.source_kernel_id)
            peer_weights.append(w)
        all_confidences = [own_weight] + peer_weights
        min_len = min(len(t) for t in all_trajectories if t)

        if min_len == 0:
            return IntegrationResult(
                integrated_trajectory=own_trajectory,
                n_contributors=1,
                mean_confidence=own_confidence,
                max_divergence=0.0,
            )

        integrated: list[Basin] = []
        max_divergence = 0.0

        for step in range(min_len):
            # Collect basins at this step from all contributors
            step_basins = [
                to_simplex(t[step]) for t in all_trajectories if len(t) > step
            ]
            step_weights = [
                c
                for c, t in zip(all_confidences, all_trajectories, strict=False)
                if len(t) > step
            ]

            if len(step_basins) == 1:
                integrated.append(step_basins[0])
                continue

            # Weighted Fréchet mean on Δ⁶³
            weight_arr = np.array(step_weights, dtype=np.float64)
            weight_arr = weight_arr / weight_arr.sum()  # normalize

            merged = frechet_mean(step_basins, weights=weight_arr)
            integrated.append(merged)

            # Track divergence: max FR distance between own and any received at this step
            own_at_step = (
                to_simplex(own_trajectory[step]) if step < len(own_trajectory) else None
            )
            if own_at_step is not None:
                for msg in received:
                    if step < len(msg.trajectory):
                        div = fisher_rao_distance(
                            own_at_step, to_simplex(msg.trajectory[step])
                        )
                        max_divergence = max(max_divergence, float(div))

        mean_conf = float(np.mean(all_confidences))

        return IntegrationResult(
            integrated_trajectory=integrated,
            n_contributors=len(all_trajectories),
            mean_confidence=mean_conf,
            max_divergence=max_divergence,
        )

    def get_state(self) -> dict[str, Any]:
        """Serializable state summary."""
        return {
            "broadcast_pending": len(self._broadcast),
            "kernel_queues": {kid: len(q) for kid, q in self._queues.items() if q},
            "total_pending": (
                len(self._broadcast) + sum(len(q) for q in self._queues.values())
            ),
        }
