"""
MetricsCollector: In-memory metrics storage for gossip learning monitoring.

Collects and stores metrics from GossipNode instances for MCP analysis.
"""

import time
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any
from datetime import datetime
import threading


@dataclass
class TrainingMetric:
    """Single training round metric."""
    timestamp: float
    round_number: int
    loss: float
    accuracy: float
    samples_trained: int


@dataclass
class ModelExchange:
    """Model exchange record."""
    timestamp: float
    round_number: int
    direction: str  # "sent" or "received"
    peer_id: str
    model_size_bytes: int
    loss: Optional[float] = None
    accuracy: Optional[float] = None


@dataclass
class PeerEvent:
    """Peer discovery/disconnect event."""
    timestamp: float
    event_type: str  # "discovered" or "disconnected"
    peer_id: str
    domain: Optional[str] = None


@dataclass
class AggregationEvent:
    """Model aggregation event."""
    timestamp: float
    round_number: int
    num_models: int
    peer_ids: List[str]
    total_samples: int


@dataclass
class NodeMetrics:
    """All metrics for a single node."""
    node_id: str
    domain: str
    start_time: float = field(default_factory=time.time)
    
    training_history: List[TrainingMetric] = field(default_factory=list)
    exchanges: List[ModelExchange] = field(default_factory=list)
    peer_events: List[PeerEvent] = field(default_factory=list)
    aggregations: List[AggregationEvent] = field(default_factory=list)
    
    current_round: int = 0
    is_running: bool = False
    final_accuracy: Optional[float] = None
    
    def to_dict(self) -> dict:
        """Convert to dictionary for JSON serialization."""
        return {
            "node_id": self.node_id,
            "domain": self.domain,
            "start_time": self.start_time,
            "current_round": self.current_round,
            "is_running": self.is_running,
            "final_accuracy": self.final_accuracy,
            "training_count": len(self.training_history),
            "exchange_count": len(self.exchanges),
            "peer_event_count": len(self.peer_events),
            "aggregation_count": len(self.aggregations),
        }


class MetricsCollector:
    """
    Collects metrics from GossipNode instances.
    
    Thread-safe metrics storage that can be attached to one or more
    GossipNode instances to collect training, exchange, and aggregation data.
    
    Example:
        collector = MetricsCollector()
        collector.attach(node1)
        collector.attach(node2)
        
        # Get all nodes' data
        data = collector.get_all_metrics()
    """
    
    def __init__(self):
        """Initialize the metrics collector."""
        self._nodes: Dict[str, NodeMetrics] = {}
        self._lock = threading.RLock()
        self._attached_nodes: Dict[str, Any] = {}  # node_id -> GossipNode
    
    def attach(self, gossip_node: Any) -> None:
        """
        Attach to a GossipNode to collect metrics.
        
        Args:
            gossip_node: A GossipNode instance to monitor
        """
        node_id = gossip_node.node_id
        domain = gossip_node.domain
        
        with self._lock:
            # Initialize node metrics
            self._nodes[node_id] = NodeMetrics(
                node_id=node_id,
                domain=domain
            )
            self._attached_nodes[node_id] = gossip_node
        
        # Register hooks on the aggregator
        aggregator = gossip_node.gl_node.aggregator

        # Training hook
        def on_training_complete(result):
            self._record_training(
                node_id=node_id,
                round_number=aggregator.current_round,
                loss=result.loss,
                accuracy=result.accuracy,
                samples=result.samples_trained
            )

        aggregator.register_hook("on_training_complete", on_training_complete)

        # Model sent hook
        def on_model_sent(peer_ids, model_size):
            for peer_id in peer_ids:
                self._record_exchange(
                    node_id=node_id,
                    round_number=aggregator.current_round,
                    direction="sent",
                    peer_id=peer_id,
                    model_size=model_size
                )

        aggregator.register_hook("on_model_sent", on_model_sent)

        # Aggregation hook
        def on_aggregation_complete(aggregated_model):
            peer_ids = [u.sender_id for u in aggregated_model.updates]
            total_samples = sum(u.sample_count for u in aggregated_model.updates)
            self._record_aggregation(
                node_id=node_id,
                round_number=aggregator.current_round,
                num_models=len(aggregated_model.updates),
                peer_ids=peer_ids,
                total_samples=total_samples
            )

        aggregator.register_hook("on_aggregation_complete", on_aggregation_complete)
    
    def detach(self, node_id: str) -> None:
        """Detach from a node (stop collecting metrics)."""
        with self._lock:
            self._attached_nodes.pop(node_id, None)
            # Note: metrics are kept for analysis
    
    def _record_training(
        self,
        node_id: str,
        round_number: int,
        loss: float,
        accuracy: float,
        samples: int
    ) -> None:
        """Record a training metric."""
        with self._lock:
            if node_id in self._nodes:
                self._nodes[node_id].training_history.append(
                    TrainingMetric(
                        timestamp=time.time(),
                        round_number=round_number,
                        loss=loss,
                        accuracy=accuracy,
                        samples_trained=samples
                    )
                )
                self._nodes[node_id].current_round = round_number
    
    def _record_exchange(
        self,
        node_id: str,
        round_number: int,
        direction: str,
        peer_id: str,
        model_size: int,
        loss: Optional[float] = None,
        accuracy: Optional[float] = None
    ) -> None:
        """Record a model exchange."""
        with self._lock:
            if node_id in self._nodes:
                self._nodes[node_id].exchanges.append(
                    ModelExchange(
                        timestamp=time.time(),
                        round_number=round_number,
                        direction=direction,
                        peer_id=peer_id,
                        model_size_bytes=model_size,
                        loss=loss,
                        accuracy=accuracy
                    )
                )
    
    def record_model_received(
        self,
        node_id: str,
        round_number: int,
        peer_id: str,
        model_size: int,
        loss: Optional[float] = None,
        accuracy: Optional[float] = None
    ) -> None:
        """Public method to record received model (called from community)."""
        self._record_exchange(
            node_id=node_id,
            round_number=round_number,
            direction="received",
            peer_id=peer_id,
            model_size=model_size,
            loss=loss,
            accuracy=accuracy
        )
    
    def record_peer_event(
        self,
        node_id: str,
        event_type: str,
        peer_id: str,
        domain: Optional[str] = None
    ) -> None:
        """Record a peer discovery or disconnect event."""
        with self._lock:
            if node_id in self._nodes:
                self._nodes[node_id].peer_events.append(
                    PeerEvent(
                        timestamp=time.time(),
                        event_type=event_type,
                        peer_id=peer_id,
                        domain=domain
                    )
                )
    
    def _record_aggregation(
        self,
        node_id: str,
        round_number: int,
        num_models: int,
        peer_ids: List[str],
        total_samples: int
    ) -> None:
        """Record an aggregation event."""
        with self._lock:
            if node_id in self._nodes:
                self._nodes[node_id].aggregations.append(
                    AggregationEvent(
                        timestamp=time.time(),
                        round_number=round_number,
                        num_models=num_models,
                        peer_ids=peer_ids,
                        total_samples=total_samples
                    )
                )
    
    def set_final_accuracy(self, node_id: str, accuracy: float) -> None:
        """Set the final test accuracy for a node."""
        with self._lock:
            if node_id in self._nodes:
                self._nodes[node_id].final_accuracy = accuracy
                self._nodes[node_id].is_running = False
    
    def set_running(self, node_id: str, running: bool) -> None:
        """Set the running state for a node."""
        with self._lock:
            if node_id in self._nodes:
                self._nodes[node_id].is_running = running
    
    # =========================================================================
    # Query Methods (used by MCP tools)
    # =========================================================================
    
    def get_all_nodes(self) -> List[str]:
        """Get list of all monitored node IDs."""
        with self._lock:
            return list(self._nodes.keys())
    
    def get_node_metrics(self, node_id: str) -> Optional[NodeMetrics]:
        """Get metrics for a specific node."""
        with self._lock:
            return self._nodes.get(node_id)
    
    def get_all_metrics(self) -> Dict[str, dict]:
        """Get summary metrics for all nodes."""
        with self._lock:
            return {
                node_id: metrics.to_dict()
                for node_id, metrics in self._nodes.items()
            }
    
    def get_accuracy_history(self, node_id: str) -> List[dict]:
        """Get accuracy history for a node."""
        with self._lock:
            if node_id not in self._nodes:
                return []
            return [
                {
                    "round": m.round_number,
                    "accuracy": m.accuracy,
                    "loss": m.loss,
                    "timestamp": m.timestamp
                }
                for m in self._nodes[node_id].training_history
            ]
    
    def get_all_accuracy_histories(self) -> Dict[str, List[dict]]:
        """Get accuracy history for all nodes."""
        with self._lock:
            return {
                node_id: self.get_accuracy_history(node_id)
                for node_id in self._nodes.keys()
            }
    
    def get_model_exchanges(self, node_id: Optional[str] = None) -> List[dict]:
        """Get model exchange records."""
        with self._lock:
            if node_id:
                if node_id not in self._nodes:
                    return []
                exchanges = self._nodes[node_id].exchanges
            else:
                exchanges = []
                for metrics in self._nodes.values():
                    for ex in metrics.exchanges:
                        exchanges.append((metrics.node_id, ex))
                
                return [
                    {
                        "node_id": nid,
                        "direction": ex.direction,
                        "peer_id": ex.peer_id,
                        "round": ex.round_number,
                        "size_kb": ex.model_size_bytes / 1024,
                        "timestamp": ex.timestamp
                    }
                    for nid, ex in exchanges
                ]
            
            return [
                {
                    "direction": ex.direction,
                    "peer_id": ex.peer_id,
                    "round": ex.round_number,
                    "size_kb": ex.model_size_bytes / 1024,
                    "timestamp": ex.timestamp
                }
                for ex in exchanges
            ]
    
    def get_network_topology(self) -> dict:
        """Get network topology based on exchanges."""
        with self._lock:
            nodes = list(self._nodes.keys())
            edges = []
            
            for node_id, metrics in self._nodes.items():
                for ex in metrics.exchanges:
                    if ex.direction == "sent":
                        edges.append({
                            "from": node_id,
                            "to": ex.peer_id,
                            "round": ex.round_number
                        })
            
            return {
                "nodes": nodes,
                "edges": edges,
                "total_exchanges": sum(
                    len(m.exchanges) for m in self._nodes.values()
                )
            }
    
    def get_convergence_analysis(self) -> dict:
        """Analyze convergence across all nodes."""
        with self._lock:
            if not self._nodes:
                return {"error": "No nodes registered"}
            
            all_accuracies = {}
            for node_id, metrics in self._nodes.items():
                if metrics.training_history:
                    all_accuracies[node_id] = [
                        m.accuracy for m in metrics.training_history
                    ]
            
            if not all_accuracies:
                return {"error": "No training data available"}
            
            # Calculate statistics
            final_accuracies = {
                nid: accs[-1] if accs else 0
                for nid, accs in all_accuracies.items()
            }
            
            initial_accuracies = {
                nid: accs[0] if accs else 0
                for nid, accs in all_accuracies.items()
            }
            
            avg_final = sum(final_accuracies.values()) / len(final_accuracies)
            avg_initial = sum(initial_accuracies.values()) / len(initial_accuracies)
            
            # Calculate improvement per round
            max_rounds = max(len(accs) for accs in all_accuracies.values())
            avg_per_round = []
            for r in range(max_rounds):
                round_accs = [
                    accs[r] for accs in all_accuracies.values()
                    if r < len(accs)
                ]
                if round_accs:
                    avg_per_round.append(sum(round_accs) / len(round_accs))
            
            return {
                "num_nodes": len(self._nodes),
                "total_rounds": max_rounds,
                "initial_avg_accuracy": avg_initial,
                "final_avg_accuracy": avg_final,
                "improvement": avg_final - avg_initial,
                "per_node_final": final_accuracies,
                "avg_accuracy_per_round": avg_per_round,
                "convergence_rate": (
                    (avg_final - avg_initial) / max_rounds
                    if max_rounds > 0 else 0
                )
            }
    
    def clear(self) -> None:
        """Clear all collected metrics."""
        with self._lock:
            self._nodes.clear()
            self._attached_nodes.clear()
    
    def export_json(self) -> dict:
        """Export all metrics as JSON-serializable dict."""
        with self._lock:
            return {
                "timestamp": time.time(),
                "nodes": {
                    node_id: {
                        **metrics.to_dict(),
                        "training_history": [
                            {
                                "round": m.round_number,
                                "loss": m.loss,
                                "accuracy": m.accuracy,
                                "samples": m.samples_trained,
                                "timestamp": m.timestamp
                            }
                            for m in metrics.training_history
                        ],
                        "exchanges": [
                            {
                                "direction": ex.direction,
                                "peer_id": ex.peer_id,
                                "round": ex.round_number,
                                "size_bytes": ex.model_size_bytes,
                                "timestamp": ex.timestamp
                            }
                            for ex in metrics.exchanges
                        ],
                        "aggregations": [
                            {
                                "round": agg.round_number,
                                "num_models": agg.num_models,
                                "peer_ids": agg.peer_ids,
                                "total_samples": agg.total_samples,
                                "timestamp": agg.timestamp
                            }
                            for agg in metrics.aggregations
                        ]
                    }
                    for node_id, metrics in self._nodes.items()
                }
            }
