"""
MeshNode — Domain-agnostic mesh cognition node.

Any application with a CfC model can participate in mesh cognition.
The node handles autonomous drift-bounded coupling — your app just
calls update_local_state() after inference and coupled_state() before.

Copyright (c) 2026 SYM.BOT Ltd. Apache 2.0 License.
"""

from typing import Optional

from .coupler import CfCMeshCoupler, CouplingConfig, CouplingDecision, CouplingDecisionType
from .metrics import KuramotoMetricsLogger


class MeshNode:
    """Domain-agnostic mesh cognition node.

    Usage:
        node = MeshNode(hidden_dim=64)

        # After each CfC inference step:
        node.update_local_state(h1, h2, confidence=0.8)

        # Add peer state (from network, another device, etc.):
        node.add_peer("peer-1", peer_h1, peer_h2, confidence=0.9)

        # Before next inference — get coupled state:
        h1, h2 = node.coupled_state()
        # Feed h1, h2 into your CfC model as hidden state inputs
    """

    def __init__(
        self,
        hidden_dim: int = 64,
        config: Optional[CouplingConfig] = None,
    ):
        self.hidden_dim = hidden_dim
        self.config = config or CouplingConfig()
        self._coupler = CfCMeshCoupler(config=self.config)
        self._metrics = KuramotoMetricsLogger()
        self._local_h1 = [0.0] * hidden_dim
        self._local_h2 = [0.0] * hidden_dim

    # -- State Exchange --

    def update_local_state(self, h1: list[float], h2: list[float], confidence: float = 0.5):
        """Update local hidden state after each CfC inference step."""
        assert len(h1) == self.hidden_dim, f"h1 dim mismatch: {len(h1)} != {self.hidden_dim}"
        assert len(h2) == self.hidden_dim, f"h2 dim mismatch: {len(h2)} != {self.hidden_dim}"
        self._local_h1 = list(h1)
        self._local_h2 = list(h2)

    def coupled_state(self) -> tuple[list[float], list[float]]:
        """Get mesh-coupled hidden state for the next inference step.

        Returns (h1, h2) — local state blended with accepted peers.
        Returns local state unchanged if no peers or all rejected.
        """
        h1, h2 = self._coupler.couple(self._local_h1, self._local_h2)

        # Log metrics
        if self._coupler.has_active_peers:
            distances = self._coupler.l2_distance(self._local_h1, self._local_h2)
            r = self._coupler.kuramoto_order_parameter(self._local_h1, self._local_h2)
            decisions = self._coupler.last_decisions
            accepted = [d for d in decisions.values() if d.decision != CouplingDecisionType.REJECTED]
            eff_alpha = sum(d.alpha for d in accepted) / len(accepted) if accepted else 0.0

            self._metrics.log(
                peer_count=self._coupler.active_peer_count,
                l2_distances=distances,
                order_parameter=r,
                effective_alpha=eff_alpha,
                decisions=decisions,
            )

        return h1, h2

    # -- Peer Management --

    def add_peer(self, agent_id: str, h1: list[float], h2: list[float], confidence: float):
        """Add or update a peer's hidden state."""
        self._coupler.update_peer(agent_id, h1, h2, confidence)

    def remove_peer(self, agent_id: str):
        self._coupler.remove_peer(agent_id)

    def remove_all_peers(self):
        self._coupler.remove_all_peers()

    # -- Metrics --

    @property
    def active_peer_count(self) -> int:
        return self._coupler.active_peer_count

    @property
    def has_active_peers(self) -> bool:
        return self._coupler.has_active_peers

    @property
    def coupling_decisions(self) -> dict[str, CouplingDecision]:
        return self._coupler.last_decisions

    @property
    def kuramoto_order_parameter(self) -> Optional[float]:
        return self._coupler.kuramoto_order_parameter(self._local_h1, self._local_h2)

    @property
    def l2_distance(self) -> Optional[tuple[float, float]]:
        return self._coupler.l2_distance(self._local_h1, self._local_h2)

    def export_metrics(self, path: str = "kuramoto_metrics.jsonl") -> str:
        """Export Kuramoto metrics to JSONL file."""
        return str(self._metrics.export(path))

    @property
    def metrics_count(self) -> int:
        return self._metrics.entry_count
