# Copyright 2026 Ali Seyhan, Baki Turhan
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Gossip Learning Node with IPv8 Integration and Fallback Support

Combines LearningNode with GossipLearningCommunity for full P2P gossip learning.
Automatically falls back to tunnel relay when IPv8 P2P fails.
"""

import asyncio
import logging
import time
from typing import Optional, Any, Callable, Dict
from enum import Enum

from quinkgl.core.learning_node import LearningNode
from quinkgl.models.base import ModelWrapper, TrainingConfig
from quinkgl.topology.base import TopologyStrategy, PeerInfo
from quinkgl.aggregation.base import AggregationStrategy
from quinkgl.network.ipv8_manager import IPv8Manager
from quinkgl.network.gossip_community import generate_community_id
from quinkgl.network.gossip_community import GossipLearningCommunity

logger = logging.getLogger(__name__)


class ConnectionMode(Enum):
    """Connection mode for the node."""
    IPV8_P2P = "ipv8_p2p"       # Direct P2P via IPv8
    TUNNEL_RELAY = "tunnel"     # Tunnel relay fallback


class GossipNode:
    """
    Complete Gossip Learning Node with IPv8 networking and tunnel fallback.

    This class combines the LearningNode framework with IPv8 P2P networking
    for full decentralized gossip learning. Falls back to tunnel relay
    when IPv8 P2P is unavailable.

    Example:
        ```python
        from quinkgl.network import GossipNode
        from quinkgl.models import PyTorchModel

        # Wrap your model
        model = PyTorchModel(my_pytorch_model)

        # Create node with tunnel fallback
        node = GossipNode(
            node_id="alice",
            domain="health",
            model=model,
            port=7000,
            tunnel_server="tunnel.example.com:50051"  # Optional fallback
        )

        # Start and run
        await node.start()
        await node.run_continuous(training_data)
        ```
    """

    def __init__(
        self,
        node_id: str,
        domain: str,
        model: ModelWrapper,
        port: int = 0,
        topology: Optional[TopologyStrategy] = None,
        aggregation: Optional[AggregationStrategy] = None,
        gossip_interval: float = 60.0,
        training_config: Optional[TrainingConfig] = None,
        auto_discovery: bool = True,
        tunnel_server: Optional[str] = None,
        enable_fallback: bool = True,
        fallback_timeout: float = 30.0
    ):
        """
        Initialize GossipNode.

        Args:
            node_id: Unique identifier for this node
            domain: Domain identifier (e.g., "health", "agriculture")
            model: Wrapped model (PyTorchModel, TensorFlowModel, or custom)
            port: UDP port for IPv8 (0 for random)
            topology: Topology strategy (defaults to RandomTopology)
            aggregation: Aggregation strategy (defaults to FedAvg)
            gossip_interval: Seconds between gossip rounds
            training_config: Configuration for local training
            auto_discovery: Enable automatic peer discovery
            tunnel_server: Tunnel server address (host:port) for fallback
            enable_fallback: Enable tunnel fallback when IPv8 fails
            fallback_timeout: Seconds to wait for IPv8 before fallback
        """
        self.node_id = node_id
        self.domain = domain
        self.port = port
        self.tunnel_server = tunnel_server
        self.enable_fallback = enable_fallback
        self.fallback_timeout = fallback_timeout

        # Get data schema hash from model
        self.data_schema_hash = model.get_data_schema_hash()

        # Import default strategies if not provided
        if topology is None:
            from quinkgl.topology import RandomTopology
            topology = RandomTopology()

        if aggregation is None:
            from quinkgl.aggregation import FedAvg
            aggregation = FedAvg()

        # Create the framework LearningNode
        self.gl_node = LearningNode(
            peer_id=node_id,
            domain=domain,
            model=model,
            topology=topology,
            aggregation=aggregation,
            data_schema_hash=self.data_schema_hash,
            gossip_interval=gossip_interval,
            training_config=training_config
        )

        # IPv8 manager
        self.ipv8_manager = IPv8Manager(node_id=node_id, port=port)
        self.ipv8_manager.domain = domain
        self.ipv8_manager.data_schema_hash = self.data_schema_hash

        # Community (set after IPv8 starts)
        self.community: Optional[GossipLearningCommunity] = None

        # Tunnel client (lazy loaded)
        self.tunnel_client = None
        self._tunnel_connected = False

        # Remote peers via tunnel
        self._tunnel_peers: Dict[str, dict] = {}  # peer_id -> {node_id, domain, schema}

        # Connection mode
        self.connection_mode: ConnectionMode = ConnectionMode.IPV8_P2P

        # Auto-discovery flag
        self.auto_discovery = auto_discovery

        # State
        self.running = False
        self._ipv8_failed = False

        # Callbacks for tunnel messages
        self._on_tunnel_peer_discovered: Optional[Callable] = None
        self._on_tunnel_model_update: Optional[Callable] = None

        logger.info(
            f"GossipNode initialized: node_id={node_id}, domain={domain}, "
            f"schema={self.data_schema_hash[:8]}..., "
            f"fallback={'enabled' if enable_fallback else 'disabled'}"
        )

    async def start(self):
        """Start the node and join the P2P network."""
        if self.running:
            logger.warning("Node already running")
            return

        logger.info(f"Starting GossipNode '{self.node_id}'...")
        logger.info(f"Attempting P2P connection (timeout: {self.fallback_timeout}s)...")

        # Try IPv8 P2P first with timeout
        ipv8_started = await self._try_start_ipv8_with_timeout()

        if ipv8_started:
            self.connection_mode = ConnectionMode.IPV8_P2P
            logger.info("✅ Using IPv8 P2P mode")
        elif self.enable_fallback and self.tunnel_server:
            # IPv8 failed or timed out, try tunnel fallback
            logger.warning("IPv8 P2P failed/timeout, falling back to tunnel relay...")
            self.connection_mode = ConnectionMode.TUNNEL_RELAY
            await self._start_tunnel_fallback()
        else:
            raise RuntimeError(
                "Failed to start networking. IPv8 failed and "
                "fallback is disabled or no tunnel server configured."
            )

        # Mark LearningNode as ready (network layer connected)
        # Note: LearningNode doesn't use _joined flag anymore, this is for GossipNode's internal tracking
        self.running = True

        mode_str = "IPv8 P2P" if self.connection_mode == ConnectionMode.IPV8_P2P else "Tunnel Relay"
        logger.info(f"✅ GossipNode '{self.node_id}' started in {mode_str} mode")

    async def _try_start_ipv8_with_timeout(self) -> bool:
        """
        Try to start IPv8 P2P networking with timeout.

        Returns:
            True if IPv8 P2P connection successful, False if timeout/error
        """
        try:
            # Step 1: Start IPv8 (with timeout)
            start_time = time.time()

            try:
                await asyncio.wait_for(
                    self.ipv8_manager.start(
                        community_class=GossipLearningCommunity,
                        node_id_param=self.node_id
                    ),
                    timeout=self.fallback_timeout
                )
            except asyncio.TimeoutError:
                logger.warning(f"IPv8 start timed out after {self.fallback_timeout}s")
                return False

            # Get community reference
            self.community = self.ipv8_manager.community

            # Update community with our specific info
            self.community.domain = self.domain
            self.community.data_schema_hash = self.data_schema_hash
            self.community._instance_community_id = generate_community_id(
                self.domain, self.data_schema_hash
            )
            type(self.community).community_id = self.community._instance_community_id

            # Setup callbacks
            self._setup_ipv8_callbacks()

            # Step 2: Wait for peer discovery with remaining timeout
            elapsed = time.time() - start_time
            remaining_timeout = max(0, self.fallback_timeout - elapsed)

            if remaining_timeout > 0:
                try:
                    await asyncio.wait_for(
                        self._wait_for_peers(),
                        timeout=remaining_timeout
                    )
                except asyncio.TimeoutError:
                    logger.warning(f"Peer discovery timed out after {self.fallback_timeout}s total")
                    # IPv8 started but no peers discovered - could still work
                    # Continue in P2P mode, peers may discover later
                    logger.info("Continuing in P2P mode (peers may discover later)")

            # Sync known peers
            self._sync_known_peers()

            peer_count = self.community.get_peer_count()
            logger.info(f"IPv8 P2P ready with {peer_count} peers discovered")

            return True

        except Exception as e:
            logger.warning(f"IPv8 P2P failed: {e}")
            self._ipv8_failed = True
            return False

    async def _wait_for_peers(self, min_peers: int = 0) -> None:
        """
        Wait for peer discovery.

        Args:
            min_peers: Minimum number of peers required (0 = any is fine)
        """
        # Poll for peer discovery with exponential backoff
        check_interval = 0.5
        max_interval = 2.0

        while True:
            peer_count = self.community.get_peer_count()

            if peer_count >= min_peers:
                logger.debug(f"Discovered {peer_count} peers")
                return

            # Wait before next check
            await asyncio.sleep(check_interval)
            check_interval = min(check_interval * 1.5, max_interval)

    def _setup_ipv8_callbacks(self):
        """Setup callbacks between IPv8 community and LearningNode."""
        # When model update received
        async def on_model_update(
            sender_id: str,
            weights: Any,
            sample_count: int,
            round_number: int,
            loss: float,
            accuracy: float
        ):
            from quinkgl.gossip.protocol import ModelUpdateMessage

            message = ModelUpdateMessage.create(
                sender_id=sender_id,
                weights=weights,
                sample_count=sample_count,
                loss=loss,
                accuracy=accuracy,
                round_number=round_number
            )

            await self.gl_node._handle_network_message(message)
            logger.debug(f"Processed model update from {sender_id}")

        self.community.on_model_update_callback = on_model_update

        # When peer discovered
        async def on_peer_discovered(peer_info):
            from quinkgl.topology.base import PeerInfo as FrameworkPeerInfo

            framework_peer_info = FrameworkPeerInfo(
                peer_id=peer_info.node_id,
                domain=peer_info.domain,
                data_schema_hash=peer_info.data_schema_hash,
                model_version=peer_info.model_version
            )

            self.gl_node.aggregator.add_peer(framework_peer_info)
            logger.info(f"Added peer {peer_info.node_id} to aggregator")

        self.community.on_peer_discovered_callback = on_peer_discovered

        # When peer leaves
        async def on_peer_left(node_id: str):
            await self.gl_node.aggregator.remove_peer(node_id)
            logger.info(f"Removed peer {node_id} from aggregator")

        self.community.on_peer_left_callback = on_peer_left

    async def _start_tunnel_fallback(self):
        """Start tunnel relay fallback."""
        try:
            from quinkgl.network.fallback import TunnelClient

            self.tunnel_client = TunnelClient(self.tunnel_server, self.node_id)
            await self.tunnel_client.connect()
            self._tunnel_connected = True

            # Setup tunnel message handlers
            self._setup_tunnel_callbacks()

            logger.info(f"✅ Tunnel fallback connected to {self.tunnel_server}")

        except Exception as e:
            logger.error(f"Tunnel fallback failed: {e}")
            raise

    def _setup_tunnel_callbacks(self):
        """Setup callbacks for tunnel messages."""
        if not self.tunnel_client:
            return

        # Handle incoming chat messages (may contain model updates)
        async def on_tunnel_message(msg):
            # Parse message - could be model update or peer discovery
            try:
                import json

                data = json.loads(msg.text)

                if data.get("type") == "MODEL_UPDATE":
                    # Handle model update via tunnel
                    if self._on_tunnel_model_update:
                        await self._on_tunnel_model_update(data)
                elif data.get("type") == "PEER_ANNOUNCE":
                    # Peer discovery announcement
                    peer_info = data.get("peer_info", {})
                    self._tunnel_peers[peer_info.get("node_id")] = peer_info

                    if self._on_tunnel_peer_discovered:
                        await self._on_tunnel_peer_discovered(peer_info)

                    # Add to aggregator
                    from quinkgl.topology.base import PeerInfo as FrameworkPeerInfo

                    framework_peer_info = FrameworkPeerInfo(
                        peer_id=peer_info.get("node_id"),
                        domain=peer_info.get("domain"),
                        data_schema_hash=peer_info.get("data_schema_hash"),
                        model_version=peer_info.get("model_version", "0.1.0")
                    )

                    self.gl_node.aggregator.add_peer(framework_peer_info)
                    logger.info(f"Added tunnel peer {peer_info.get('node_id')} to aggregator")

            except json.JSONDecodeError:
                # Regular chat message, ignore
                pass

        self.tunnel_client.on_chat_message = on_tunnel_message

        # Handle peer list updates from tunnel
        async def on_peer_list(peer_ids: list):
            logger.info(f"Tunnel server reports {len(peer_ids)} connected peers")

        self.tunnel_client.on_peer_list = on_peer_list

    def _sync_known_peers(self):
        """Sync known peers from community to LearningNode."""
        if not self.community:
            return

        from quinkgl.topology.base import PeerInfo as FrameworkPeerInfo

        connected_peer_ids = []
        for peer_info in self.community.get_compatible_peers():
            framework_peer_info = FrameworkPeerInfo(
                peer_id=peer_info.node_id,
                domain=peer_info.domain,
                data_schema_hash=peer_info.data_schema_hash,
                model_version=peer_info.model_version
            )
            self.gl_node.aggregator.add_peer(framework_peer_info)
            connected_peer_ids.append(peer_info.node_id)

        logger.info(f"Synced {self.community.get_peer_count()} peers to aggregator")

    async def run_continuous(self, data=None, data_provider=None):
        """
        Run continuous gossip learning.

        Args:
            data: Training data (single dataset)
            data_provider: Callable that returns training data per round
        """
        if not self.running:
            raise RuntimeError("Node must be started before running")

        # Setup send callback based on connection mode
        if self.connection_mode == ConnectionMode.IPV8_P2P:
            async def send_to_peer(peer_id: str, message):
                """Send message via IPv8."""
                self.community.send_model_update(
                    target_node_id=peer_id,
                    weights=message.weights,
                    sample_count=message.sample_count,
                    round_number=message.round_number,
                    loss=message.loss,
                    accuracy=message.accuracy
                )
        else:
            async def send_to_peer(peer_id: str, message):
                """Send message via tunnel fallback."""
                await self._send_model_update_via_tunnel(peer_id, message)

        self.gl_node.aggregator.send_message_callback = send_to_peer

        # Sync peers before starting
        if self.connection_mode == ConnectionMode.IPV8_P2P:
            self._sync_known_peers()

        # Announce ourselves to other tunnel peers
        if self.connection_mode == ConnectionMode.TUNNEL_RELAY:
            await self._announce_to_tunnel()

        # Run the gossip loop
        await self.gl_node.run_continuous(data_provider=data_provider or data)

    async def _send_model_update_via_tunnel(self, peer_id: str, message):
        """Send model update via tunnel relay."""
        if not self.tunnel_client or not self._tunnel_connected:
            logger.warning(f"Cannot send to {peer_id}: tunnel not connected")
            return

        import json
        from quinkgl.network.model_serializer import serialize_model

        try:
            # Serialize weights
            weights_bytes = serialize_model(self.node_id, message.weights)

            payload = {
                "type": "MODEL_UPDATE",
                "sender_id": self.node_id,
                "domain": self.domain,
                "data_schema_hash": self.data_schema_hash,
                "round_number": message.round_number,
                "sample_count": message.sample_count,
                "loss": message.loss,
                "accuracy": message.accuracy,
                "weights": weights_bytes.hex()  # Convert bytes to hex for JSON
            }

            # Send via tunnel chat (hacky but works for fallback)
            await self.tunnel_client.send_chat_message(peer_id, json.dumps(payload))
            logger.debug(f"Sent model update to {peer_id} via tunnel")

        except Exception as e:
            logger.error(f"Failed to send model update via tunnel: {e}")

    async def _announce_to_tunnel(self):
        """Announce ourselves to other peers via tunnel."""
        if not self.tunnel_client:
            return

        import json

        announcement = {
            "type": "PEER_ANNOUNCE",
            "peer_info": {
                "node_id": self.node_id,
                "domain": self.domain,
                "data_schema_hash": self.data_schema_hash,
                "model_version": "0.1.0"
            }
        }

        # Broadcast to all known peers
        for peer_id in self._tunnel_peers.keys():
            try:
                await self.tunnel_client.send_chat_message(peer_id, json.dumps(announcement))
            except Exception as e:
                logger.debug(f"Failed to announce to {peer_id}: {e}")

    def stop(self):
        """Stop the node."""
        if not self.running:
            return

        logger.info(f"Stopping GossipNode '{self.node_id}'...")

        self.gl_node.stop()
        self.running = False

    async def shutdown(self):
        """Full shutdown including IPv8 and tunnel."""
        self.stop()

        # Stop IPv8
        await self.ipv8_manager.stop()

        # Stop tunnel client
        if self.tunnel_client and self._tunnel_connected:
            try:
                await self.tunnel_client.close()
                logger.info("Tunnel client disconnected")
            except Exception as e:
                logger.warning(f"Error closing tunnel client: {e}")

        logger.info(f"GossipNode '{self.node_id}' shutdown complete")

    def get_stats(self) -> dict:
        """Get node statistics."""
        ipv8_stats = self.ipv8_manager.get_stats()

        return {
            "node_id": self.node_id,
            "domain": self.domain,
            "data_schema_hash": self.data_schema_hash,
            "running": self.running,
            "current_round": self.gl_node.current_round,
            "connection_mode": self.connection_mode.value,
            "ipv8_peers": self.community.get_peer_count() if self.community else 0,
            "tunnel_peers": len(self._tunnel_peers),
            "ipv8_port": ipv8_stats.get("port"),
            "tunnel_server": self.tunnel_server,
            "known_peers": [
                p.node_id for p in self.community.get_compatible_peers()
            ] if self.community else list(self._tunnel_peers.keys())
        }

    def get_model(self) -> ModelWrapper:
        """Get the underlying model wrapper."""
        return self.gl_node.get_model()

    def is_using_fallback(self) -> bool:
        """Check if node is using tunnel fallback."""
        return self.connection_mode == ConnectionMode.TUNNEL_RELAY

    def get_connection_mode(self) -> ConnectionMode:
        """Get current connection mode."""
        return self.connection_mode
