# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
Multi-Device Mesh Network

Architecture:
- One GATEWAY (usually the Pi) - runs 24/7, owns the brain
- Multiple NODES (laptop, phone, other devices) - connect to gateway

Gateway:
- Runs the LLM, memory, and skills
- Accepts connections from nodes via WebSocket
- Routes messages between nodes
- Persists state

Nodes:
- Lightweight - just a connection to gateway
- Can expose local tools (clipboard, notifications, file system)
- Can run headless or with UI

Protocol:
- WebSocket for real-time communication
- JSON messages with type, payload, metadata
- Heartbeat for connection health
"""

import asyncio
import json
import logging
import secrets
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Callable

logger = logging.getLogger(__name__)

# Try to import websockets for exception handling
# The actual websockets module is imported locally in methods that need it
try:
    from websockets.exceptions import ConnectionClosed as WebSocketConnectionClosed
except ImportError:
    # Create a dummy exception that will never be raised if websockets isn't installed
    class WebSocketConnectionClosed(Exception):
        pass


# Use centralized paths
try:
    from ..paths import DATA_DIR, MESH_DIR

    MESH_CONFIG_FILE = MESH_DIR / "mesh_config.json"
except ImportError:
    DATA_DIR = Path.home() / ".familiar" / "data"
    MESH_DIR = DATA_DIR / "mesh"
    MESH_CONFIG_FILE = MESH_DIR / "mesh_config.json"


# Message types
class MsgType:
    # Connection
    HELLO = "hello"
    AUTH = "auth"
    AUTH_OK = "auth_ok"
    AUTH_FAIL = "auth_fail"
    HEARTBEAT = "heartbeat"
    DISCONNECT = "disconnect"

    # Chat
    USER_MESSAGE = "user_message"
    ASSISTANT_RESPONSE = "assistant_response"
    TYPING = "typing"

    # Tools
    TOOL_REQUEST = "tool_request"
    TOOL_RESPONSE = "tool_response"
    TOOL_REGISTER = "tool_register"

    # State sync
    SYNC_REQUEST = "sync_request"
    SYNC_RESPONSE = "sync_response"

    # Admin
    NODE_LIST = "node_list"
    NODE_STATUS = "node_status"

    # Peer gateway (mesh Phase 1)
    PEER_HELLO = "peer_hello"
    PEER_HELLO_ACK = "peer_hello_ack"
    PEER_TOOLS = "peer_tools"


@dataclass
class Message:
    """Message format for mesh communication."""

    type: str
    payload: dict
    node_id: str = ""
    timestamp: str = ""
    msg_id: str = ""

    def __post_init__(self):
        if not self.timestamp:
            self.timestamp = datetime.now().isoformat()
        if not self.msg_id:
            self.msg_id = secrets.token_hex(8)

    def to_json(self) -> str:
        return json.dumps(asdict(self))

    @classmethod
    def from_json(cls, data: str) -> "Message":
        d = json.loads(data)
        return cls(**d)


@dataclass
class NodeInfo:
    """Information about a connected node."""

    node_id: str
    name: str
    type: str  # "gateway", "desktop", "mobile", "server"
    connected_at: str
    last_seen: str
    tools: list  # Tools this node provides
    capabilities: list  # What it can do
    address: str = ""


class MeshConfig:
    """Mesh network configuration."""

    def __init__(self):
        self.config = self._load()

    def _load(self) -> dict:
        if MESH_CONFIG_FILE.exists():
            try:
                return json.loads(MESH_CONFIG_FILE.read_text())
            except (json.JSONDecodeError, IOError, OSError) as e:
                logger.warning(f"Failed to load mesh config: {e}")
        return {
            "node_id": secrets.token_hex(8),
            "node_name": "familiar",
            "node_type": "gateway",
            "secret_key": secrets.token_hex(32),
            "gateway_host": "127.0.0.1",
            "gateway_port": 18789,
            "allowed_nodes": [],  # Empty = allow any with correct key
            "auto_discovery": True,
            # Phase 1: Discovery
            "discovery_enabled": False,
            "discovery_mode": "mdns",  # 'mdns', 'manual', 'both'
            "auto_connect_trusted": True,  # Auto-connect to known trusted peers
            "peer_gateways": [],  # Manual: ["host:port", ...]
            "max_peers": 16,
        }

    def save(self):
        MESH_DIR.mkdir(parents=True, exist_ok=True)
        MESH_CONFIG_FILE.write_text(json.dumps(self.config, indent=2))
        try:
            MESH_CONFIG_FILE.chmod(0o600)
        except OSError:
            logger.warning("Could not set permissions on mesh config file")

    def get(self, key: str, default=None):
        return self.config.get(key, default)

    def set(self, key: str, value):
        self.config[key] = value
        self.save()


# ============================================================
# GATEWAY - The central hub
# ============================================================


class MeshGateway:
    """
    Central gateway that nodes connect to.
    Runs on the Pi (or main server).

    Usage:
        gateway = MeshGateway(agent)
        await gateway.start()
    """

    def __init__(self, agent, host: str = "0.0.0.0", port: int = 18789):
        self.agent = agent
        self.host = host
        self.port = port
        self.config = MeshConfig()
        self.nodes: dict[str, NodeInfo] = {}
        self.connections: dict[str, Any] = {}  # node_id -> websocket
        self.node_tools: dict[str, list] = {}  # node_id -> tools
        self._running = False
        self._server = None

        # Phase 1: Peer gateway connections
        self.peer_gateways: dict[str, Any] = {}  # node_id -> MeshNode
        self.peer_gateway_info: dict[str, dict] = {}  # node_id -> peer metadata
        self.peer_gateway_tools: dict[str, list] = {}  # node_id -> tools
        self.discovery = None  # MeshDiscovery instance

        # Phase 2: Trust manager
        self.trust_manager = None  # MeshTrustManager instance
        try:
            from .trust import MeshTrustManager

            self.trust_manager = MeshTrustManager(mesh_dir=MESH_DIR)
        except Exception as e:
            logger.debug(f"Trust manager not available: {e}")

        # Phase 3: Memory bridge
        self.memory_bridge = None  # MemoryBridge instance

    async def start(self):
        """Start the gateway server."""
        try:
            import websockets
        except ImportError:
            raise RuntimeError("websockets not available. Check server logs.")

        self._running = True

        logger.info(f"Starting mesh gateway on {self.host}:{self.port}")

        if self.host == "0.0.0.0":
            logger.warning(
                "Mesh gateway bound to all interfaces (0.0.0.0). "
                "Use --gateway-host 127.0.0.1 to restrict to localhost."
            )

        self._server = await websockets.serve(self._handle_connection, self.host, self.port)

        # Start heartbeat checker
        asyncio.create_task(self._heartbeat_loop())

        logger.info(f"Mesh gateway running. Secret key: {self.config.get('secret_key')[:8]}...")

        # Phase 1: Start discovery if enabled
        if self.config.get("discovery_enabled", False):
            try:
                from .discovery import MeshDiscovery

                # Gather skill names for advertisement
                skill_names = []
                try:
                    from ..tools import get_tool_registry

                    skill_names = [t.get("name", "") for t in get_tool_registry().get_schemas()]
                except Exception:
                    pass

                self.discovery = MeshDiscovery(
                    node_id=self.config.get("node_id"),
                    node_name=self.config.get("node_name", "familiar"),
                    node_type=self.config.get("node_type", "gateway"),
                    port=self.port,
                    skills=skill_names,
                    fingerprint=self.trust_manager.get_our_fingerprint()
                    if self.trust_manager
                    else "",
                    mesh_dir=MESH_DIR,
                    max_peers=self.config.get("max_peers", 16),
                    on_peer_found=self._on_peer_discovered,
                    on_peer_lost=self._on_peer_lost,
                    on_peer_updated=self._on_peer_updated,
                )
                self.discovery.start()
            except ImportError:
                logger.warning("Mesh discovery module not available")
            except Exception as e:
                logger.warning(f"Failed to start mesh discovery: {e}")

        # Phase 1: Connect to manually configured peer gateways
        for peer_addr in self.config.get("peer_gateways", []):
            try:
                host, port_str = peer_addr.rsplit(":", 1)
                asyncio.create_task(self.connect_to_peer_gateway(host, int(port_str)))
            except (ValueError, TypeError) as e:
                logger.warning(f"Invalid peer gateway address '{peer_addr}': {e}")

    async def stop(self):
        """Stop the gateway."""
        self._running = False

        # Phase 1: Stop discovery
        if self.discovery:
            try:
                self.discovery.stop()
            except Exception as e:
                logger.debug(f"Error stopping discovery: {e}")

        # Phase 1: Disconnect peer gateways
        for peer_id, peer_node in list(self.peer_gateways.items()):
            try:
                await peer_node.disconnect()
            except Exception as e:
                logger.debug(f"Error disconnecting peer gateway {peer_id}: {e}")
        self.peer_gateways.clear()
        self.peer_gateway_info.clear()
        self.peer_gateway_tools.clear()

        if self._server:
            self._server.close()
            await self._server.wait_closed()

        # Disconnect all nodes
        for ws in self.connections.values():
            await ws.close()

    async def _handle_connection(self, websocket, path):
        """Handle a new node connection."""
        node_id = None

        try:
            # Wait for hello/auth
            raw = await asyncio.wait_for(websocket.recv(), timeout=10)
            msg = Message.from_json(raw)

            if msg.type != MsgType.HELLO and msg.type != MsgType.PEER_HELLO:
                await websocket.close(1002, "Expected HELLO or PEER_HELLO")
                return

            is_peer_gateway = msg.type == MsgType.PEER_HELLO

            # Verify auth
            provided_key = msg.payload.get("secret_key", "")
            if provided_key != self.config.get("secret_key"):
                await websocket.send(
                    Message(
                        type=MsgType.AUTH_FAIL, payload={"reason": "Invalid secret key"}
                    ).to_json()
                )
                await websocket.close(1002, "Auth failed")
                return

            # Register node
            node_id = msg.payload.get("node_id", secrets.token_hex(8))

            if is_peer_gateway:
                # Peer gateway: register as gateway with routing capability
                node_info = NodeInfo(
                    node_id=node_id,
                    name=msg.payload.get("name", "unknown"),
                    type="gateway",
                    connected_at=datetime.now().isoformat(),
                    last_seen=datetime.now().isoformat(),
                    tools=msg.payload.get("tools", []),
                    capabilities=["routing"] + msg.payload.get("capabilities", []),
                    address=str(websocket.remote_address),
                )

                self.peer_gateways[node_id] = websocket
                self.peer_gateway_info[node_id] = {
                    "name": node_info.name,
                    "address": str(websocket.remote_address),
                    "connected_at": node_info.connected_at,
                }
                self.peer_gateway_tools[node_id] = msg.payload.get("tools", [])
                self.nodes[node_id] = node_info
                self.connections[node_id] = websocket

                # Send auth OK with our routing table
                local_node_ids = [nid for nid in self.connections if nid != node_id]
                await websocket.send(
                    Message(
                        type=MsgType.AUTH_OK,
                        payload={
                            "gateway_id": self.config.get("node_id"),
                            "gateway_name": self.config.get("node_name"),
                            "your_node_id": node_id,
                            "routing_table": local_node_ids,
                        },
                    ).to_json()
                )

                logger.info(
                    f"Peer gateway connected: {node_info.name} ({node_id}) "
                    f"with {len(node_info.tools)} tools"
                )
            else:
                # Regular node registration
                node_info = NodeInfo(
                    node_id=node_id,
                    name=msg.payload.get("name", "unknown"),
                    type=msg.payload.get("type", "unknown"),
                    connected_at=datetime.now().isoformat(),
                    last_seen=datetime.now().isoformat(),
                    tools=msg.payload.get("tools", []),
                    capabilities=msg.payload.get("capabilities", []),
                    address=str(websocket.remote_address),
                )

                self.nodes[node_id] = node_info
                self.connections[node_id] = websocket
                self.node_tools[node_id] = msg.payload.get("tools", [])

                # Send auth OK
                await websocket.send(
                    Message(
                        type=MsgType.AUTH_OK,
                        payload={
                            "gateway_id": self.config.get("node_id"),
                            "gateway_name": self.config.get("node_name"),
                            "your_node_id": node_id,
                        },
                    ).to_json()
                )

                logger.info(f"Node connected: {node_info.name} ({node_id})")

            # Handle messages
            await self._handle_node_messages(node_id, websocket)

        except asyncio.TimeoutError:
            logger.warning("Connection timeout during handshake")
        except Exception as e:
            logger.error(f"Connection error: {e}")
        finally:
            # Cleanup
            if node_id:
                self.nodes.pop(node_id, None)
                self.connections.pop(node_id, None)
                self.node_tools.pop(node_id, None)
                logger.info(f"Node disconnected: {node_id}")

    async def _handle_node_messages(self, node_id: str, websocket):
        """Handle messages from a connected node."""
        async for raw in websocket:
            try:
                msg = Message.from_json(raw)
                msg.node_id = node_id

                # Update last seen
                if node_id in self.nodes:
                    self.nodes[node_id].last_seen = datetime.now().isoformat()

                if msg.type == MsgType.HEARTBEAT:
                    await websocket.send(
                        Message(type=MsgType.HEARTBEAT, payload={"status": "ok"}).to_json()
                    )

                elif msg.type == MsgType.USER_MESSAGE:
                    # Process through agent
                    await self._handle_user_message(node_id, websocket, msg)

                elif msg.type == MsgType.TOOL_RESPONSE:
                    # Tool result from node
                    await self._handle_tool_response(node_id, msg)

                elif msg.type == MsgType.TOOL_REGISTER:
                    # Node registering new tools
                    self.node_tools[node_id] = msg.payload.get("tools", [])
                    logger.info(f"Node {node_id} registered tools: {self.node_tools[node_id]}")

                elif msg.type == MsgType.PEER_TOOLS:
                    # Peer gateway sharing its tool registry
                    self.peer_gateway_tools[node_id] = msg.payload.get("tools", [])
                    logger.info(
                        f"Peer gateway {node_id} tools updated: {len(self.peer_gateway_tools[node_id])} tools"
                    )

                elif msg.type == MsgType.NODE_LIST:
                    # Return list of connected nodes
                    await websocket.send(
                        Message(
                            type=MsgType.NODE_LIST,
                            payload={"nodes": [asdict(n) for n in self.nodes.values()]},
                        ).to_json()
                    )

            except Exception as e:
                logger.error(f"Error handling message: {e}")

    async def _handle_user_message(self, node_id: str, websocket, msg: Message):
        """Process a user message through the agent."""
        user_text = msg.payload.get("text", "")
        chat_id = msg.payload.get("chat_id", node_id)

        # Send typing indicator
        await websocket.send(Message(type=MsgType.TYPING, payload={"status": "thinking"}).to_json())

        # Run agent (in thread to not block)
        loop = asyncio.get_event_loop()
        response = await loop.run_in_executor(
            None, self.agent.chat, user_text, chat_id, f"mesh:{node_id}"
        )

        # Send response
        await websocket.send(
            Message(
                type=MsgType.ASSISTANT_RESPONSE, payload={"text": response, "chat_id": chat_id}
            ).to_json()
        )

    async def _handle_tool_response(self, node_id: str, msg: Message):
        """Handle tool response from a node."""
        # This would be used when gateway requests a tool on a node
        # Store result for pending request
        pass

    async def _heartbeat_loop(self):
        """Check node health periodically."""
        while self._running:
            await asyncio.sleep(30)

            now = datetime.now()
            dead_nodes = []

            for node_id, info in self.nodes.items():
                last_seen = datetime.fromisoformat(info.last_seen)
                if (now - last_seen).total_seconds() > 90:
                    dead_nodes.append(node_id)

            for node_id in dead_nodes:
                logger.warning(f"Node {node_id} timed out")
                ws = self.connections.pop(node_id, None)
                self.nodes.pop(node_id, None)
                if ws:
                    await ws.close()

    async def broadcast(self, msg: Message, exclude: list = None):
        """Broadcast message to all nodes."""
        exclude = exclude or []
        for node_id, ws in self.connections.items():
            if node_id not in exclude:
                try:
                    await ws.send(msg.to_json())
                except (WebSocketConnectionClosed, OSError) as e:
                    logger.debug(f"Failed to broadcast to {node_id}: {e}")

    async def send_to_node(self, node_id: str, msg: Message) -> bool:
        """Send message to specific node, routing via peer gateways if needed."""
        # Try direct connection first
        ws = self.connections.get(node_id)
        if ws:
            try:
                await ws.send(msg.to_json())
                return True
            except (WebSocketConnectionClosed, OSError) as e:
                logger.debug(f"Failed to send to {node_id}: {e}")

        # Routing fallback: forward via peer gateways
        for peer_id, peer_ws in self.peer_gateways.items():
            peer_info = self.peer_gateway_info.get(peer_id, {})
            routing_table = peer_info.get("routing_table", [])
            if node_id in routing_table:
                try:
                    await peer_ws.send(msg.to_json())
                    logger.debug(f"Routed message for {node_id} via peer gateway {peer_id}")
                    return True
                except (WebSocketConnectionClosed, OSError) as e:
                    logger.debug(f"Failed to route via {peer_id}: {e}")

        return False

    async def request_tool(
        self, node_id: str, tool_name: str, tool_input: dict, timeout: int = 30
    ) -> str:
        """Request a tool execution on a remote node."""
        ws = self.connections.get(node_id)
        if not ws:
            return f"Node {node_id} not connected"

        request_id = secrets.token_hex(8)

        await ws.send(
            Message(
                type=MsgType.TOOL_REQUEST,
                payload={"request_id": request_id, "tool": tool_name, "input": tool_input},
            ).to_json()
        )

        # Wait for response (simplified - real impl would use futures)
        # For now, return immediately
        return f"Tool request sent to {node_id}"

    def get_all_tools(self) -> list:
        """Get tools from all nodes plus local tools."""
        all_tools = []

        # Local tools
        from ..tools import get_tool_registry

        local_tools = get_tool_registry().get_schemas()
        for t in local_tools:
            t["_source"] = "local"
        all_tools.extend(local_tools)

        # Node tools
        for node_id, tools in self.node_tools.items():
            for t in tools:
                t["_source"] = f"node:{node_id}"
            all_tools.extend(tools)

        # Peer gateway tools (Phase 1)
        for peer_id, tools in self.peer_gateway_tools.items():
            for t in tools:
                t["_source"] = f"peer:{peer_id}"
            all_tools.extend(tools)

        return all_tools

    # ── Phase 1: Peer gateway connections ──

    async def connect_to_peer_gateway(
        self,
        host: str,
        port: int,
        secret_key: str = None,
    ):
        """
        Connect to another gateway as a peer (not as a subordinate node).
        Uses PEER_HELLO instead of HELLO to signal peer-level connection.
        """
        try:
            import websockets
        except ImportError:
            logger.error("websockets not installed")
            return

        if secret_key is None:
            secret_key = self.config.get("secret_key", "")

        url = f"ws://{host}:{port}"
        logger.info(f"Connecting to peer gateway: {url}")

        try:
            ws = await asyncio.wait_for(websockets.connect(url), timeout=10)

            # Send PEER_HELLO (not HELLO)
            my_tools = []
            try:
                from ..tools import get_tool_registry

                my_tools = get_tool_registry().get_schemas()
            except Exception:
                pass

            await ws.send(
                Message(
                    type=MsgType.PEER_HELLO,
                    payload={
                        "node_id": self.config.get("node_id"),
                        "name": self.config.get("node_name"),
                        "type": "gateway",
                        "secret_key": secret_key,
                        "tools": [
                            {"name": t.get("name", ""), "description": t.get("description", "")}
                            for t in my_tools
                        ],
                        "capabilities": ["chat", "tools", "peer_gateway"],
                        "version": "2.7.0",
                        # Phase 2: Include prekey bundle for trust establishment
                        "prekey_bundle": self.trust_manager.get_prekey_bundle_dict()
                        if self.trust_manager
                        else None,
                    },
                ).to_json()
            )

            # Wait for auth response
            raw = await asyncio.wait_for(ws.recv(), timeout=10)
            msg = Message.from_json(raw)

            if msg.type == MsgType.AUTH_FAIL:
                logger.error(f"Peer gateway auth failed: {msg.payload.get('reason')}")
                await ws.close()
                return

            if msg.type == MsgType.AUTH_OK:
                peer_id = msg.payload.get("gateway_id", "")
                peer_name = msg.payload.get("gateway_name", host)

                self.peer_gateway_info[peer_id] = {
                    "host": host,
                    "port": port,
                    "name": peer_name,
                    "connected_at": datetime.now().isoformat(),
                }
                self.peer_gateways[peer_id] = ws

                # Phase 2: Initiate trust handshake if peer sent a prekey bundle
                their_bundle = msg.payload.get("prekey_bundle")
                if self.trust_manager and their_bundle:
                    try:
                        init_hex = self.trust_manager.initiate_handshake(
                            peer_id, peer_name, their_bundle
                        )
                        if init_hex:
                            await ws.send(
                                Message(
                                    type=MsgType.PEER_HELLO_ACK,
                                    payload={
                                        "x3dh_init": init_hex,
                                        "identity_key": self.trust_manager.get_prekey_bundle_dict().get(
                                            "ik", ""
                                        ),
                                    },
                                ).to_json()
                            )
                    except Exception as e:
                        logger.warning(f"Trust handshake failed with {peer_name}: {e}")

                logger.info(f"Connected to peer gateway: {peer_name} ({peer_id[:8]}...)")

                # Phase 4: Auto-register RemoteAgent in orchestrator
                peer_tools = msg.payload.get("tools", [])
                self._register_remote_agent(peer_id, peer_name, peer_tools)

                # Request their tool list
                await ws.send(Message(type=MsgType.NODE_LIST, payload={}).to_json())

                # Listen for messages in background
                asyncio.create_task(self._peer_gateway_listener(peer_id, ws))

        except asyncio.TimeoutError:
            logger.warning(f"Timeout connecting to peer gateway {host}:{port}")
        except Exception as e:
            logger.warning(f"Failed to connect to peer gateway {host}:{port}: {e}")

    async def _peer_gateway_listener(self, peer_id: str, ws):
        """Listen for messages from a peer gateway."""
        try:
            async for raw in ws:
                try:
                    msg = Message.from_json(raw)

                    if msg.type == MsgType.HEARTBEAT:
                        await ws.send(
                            Message(type=MsgType.HEARTBEAT, payload={"status": "ok"}).to_json()
                        )

                    elif msg.type == MsgType.PEER_HELLO_ACK:
                        # Phase 2: Peer accepted our handshake with X3DH response
                        x3dh_init = msg.payload.get("x3dh_init")
                        their_ik = msg.payload.get("identity_key")
                        if self.trust_manager and x3dh_init and their_ik:
                            try:
                                peer_info = self.peer_gateway_info.get(peer_id, {})
                                self.trust_manager.accept_handshake(
                                    peer_id,
                                    peer_info.get("name", peer_id),
                                    x3dh_init,
                                    their_ik,
                                )
                                code = self.trust_manager.get_verification_code(peer_id)
                                if code:
                                    logger.info(
                                        f"Encrypted session with {peer_info.get('name', peer_id)}: "
                                        f"verification code {code}"
                                    )
                            except Exception as e:
                                logger.warning(f"Trust handshake accept failed: {e}")

                    elif msg.type == MsgType.PEER_TOOLS:
                        self.peer_gateway_tools[peer_id] = msg.payload.get("tools", [])

                    elif msg.type == MsgType.TOOL_RESPONSE:
                        # Response to a tool request we sent to this peer
                        await self._handle_tool_response(peer_id, msg)

                    elif msg.type == MsgType.TOOL_REQUEST:
                        # Could be a regular tool request OR a mesh event
                        if msg.payload.get("mesh_event"):
                            # Phase 3: Mesh event — dispatch to local bus handlers
                            self._handle_mesh_event(peer_id, msg.payload)
                        else:
                            # Regular tool request from peer
                            await self._handle_tool_response(peer_id, msg)

                except Exception as e:
                    logger.error(f"Error handling peer gateway message: {e}")

        except (WebSocketConnectionClosed, OSError, asyncio.CancelledError):
            pass
        finally:
            self.peer_gateways.pop(peer_id, None)
            self.peer_gateway_info.pop(peer_id, None)
            self.peer_gateway_tools.pop(peer_id, None)

            # Phase 4: Unregister remote agent and skills
            self._unregister_remote_agent(peer_id)

            logger.info(f"Peer gateway disconnected: {peer_id[:8]}...")

    async def request_peer_tool(
        self,
        peer_id: str,
        tool_name: str,
        tool_input: dict,
        timeout: int = 30,
    ) -> str:
        """Request a tool execution on a peer gateway."""
        ws = self.peer_gateways.get(peer_id)
        if not ws:
            return f"Peer gateway {peer_id} not connected"

        request_id = secrets.token_hex(8)

        await ws.send(
            Message(
                type=MsgType.TOOL_REQUEST,
                payload={
                    "request_id": request_id,
                    "tool": tool_name,
                    "input": tool_input,
                },
            ).to_json()
        )

        return f"Tool request sent to peer {peer_id}"

    def get_peer_gateways(self) -> list:
        """Get info about connected peer gateways."""
        return [
            {
                "node_id": pid,
                "name": info.get("name", ""),
                "host": info.get("host", ""),
                "port": info.get("port", 0),
                "connected_at": info.get("connected_at", ""),
                "tools": len(self.peer_gateway_tools.get(pid, [])),
            }
            for pid, info in self.peer_gateway_info.items()
        ]

    # ── Phase 1: Discovery callbacks ──

    def _on_peer_discovered(self, peer):
        """Called by MeshDiscovery when a new peer appears on the network."""
        logger.info(
            f"Discovery: new peer '{peer.node_name}' at {peer.host}:{peer.port} "
            f"(trust={peer.trust_level})"
        )
        # Auto-connect if trusted and configured to do so
        if (
            peer.trust_level == "trusted"
            and self.config.get("auto_connect_trusted", True)
            and peer.node_id not in self.peer_gateways
        ):
            logger.info(f"Auto-connecting to trusted peer: {peer.node_name}")
            asyncio.ensure_future(self.connect_to_peer_gateway(peer.host, peer.port))

    def _on_peer_lost(self, node_id: str):
        """Called by MeshDiscovery when a peer disappears."""
        logger.info(f"Discovery: peer lost {node_id[:8]}...")

    def _on_peer_updated(self, peer):
        """Called by MeshDiscovery when a peer's info changes."""
        logger.debug(f"Discovery: peer updated {peer.node_name}")

    # ── Phase 3: Mesh event handling ──

    def _handle_mesh_event(self, peer_id: str, payload: dict):
        """
        Handle a mesh event received from a peer gateway.
        Routes mesh.memory.* events to the MemoryBridge.
        """
        topic = payload.get("topic", "")
        data = payload.get("data", {})

        if not topic.startswith("mesh."):
            logger.warning(f"Ignoring non-mesh event from peer: {topic}")
            return

        if self.memory_bridge:
            if topic == "mesh.memory.query":
                response = self.memory_bridge.handle_remote_query(data)
                if response:
                    requesting_node = data.get("requesting_node", "")
                    # If the requesting node is the peer, send response back
                    if requesting_node == peer_id or not requesting_node:
                        self.memory_bridge._send_to_peer(
                            peer_id, "mesh.memory.query_response", response
                        )
            elif topic == "mesh.memory.query_response":
                self.memory_bridge.handle_remote_response(data)
            elif topic == "mesh.memory.push":
                self.memory_bridge.handle_remote_push(data)
            elif topic == "mesh.memory.revoke":
                self.memory_bridge.handle_remote_revoke(data)
            else:
                logger.debug(f"Unhandled mesh event: {topic}")

    def init_memory_bridge(self, memory, hipaa_enabled: bool = False):
        """
        Initialize the memory bridge with a Memory instance.
        Called after the agent and memory are set up.
        """
        try:
            from .memory_bridge import MemoryBridge

            self.memory_bridge = MemoryBridge(
                memory=memory,
                node_id=self.config.get("node_id"),
                node_name=self.config.get("node_name", "familiar"),
                mesh_dir=MESH_DIR,
                trust_manager=self.trust_manager,
                hipaa_enabled=hipaa_enabled,
            )
            self.memory_bridge.set_gateway(self)

            # Register bus handlers if bus is available
            try:
                from ..bus import get_skill_bus

                bus = get_skill_bus()
                self.memory_bridge.register_bus_handlers(bus)
            except Exception:
                pass

            logger.info("Memory bridge initialized")
        except Exception as e:
            logger.warning(f"Failed to init memory bridge: {e}")

    # ── Phase 4: Remote agent registration ──

    # Orchestrator reference (set externally after init)
    _orchestrator = None

    def set_orchestrator(self, orchestrator):
        """Set reference to the Orchestrator for remote agent registration."""
        self._orchestrator = orchestrator

    def _register_remote_agent(self, peer_id: str, peer_name: str, tools: list):
        """
        Register a RemoteAgent and MeshSkillEndpoints when a peer connects.
        """
        try:
            from .mesh_skill_endpoint import register_remote_skills
            from .remote_agent import RemoteAgent

            skill_names = [t.get("name", "") for t in tools if t.get("name")]

            # Register RemoteAgent in orchestrator
            if self._orchestrator:
                remote = RemoteAgent(
                    node_id=peer_id,
                    node_name=peer_name,
                    skills=skill_names,
                    gateway=self,
                    trust_manager=self.trust_manager,
                )
                self._orchestrator.registry.register(remote)
                logger.info(
                    f"RemoteAgent registered: {peer_name} ({peer_id[:8]}...) "
                    f"skills={skill_names[:5]}"
                )

            # Register skills on SkillBus
            try:
                from ..bus import get_skill_bus

                bus = get_skill_bus()
                register_remote_skills(
                    bus=bus,
                    node_id=peer_id,
                    node_name=peer_name,
                    skills=tools,
                    gateway=self,
                    trust_manager=self.trust_manager,
                )
            except Exception as e:
                logger.debug(f"SkillBus registration skipped: {e}")

        except Exception as e:
            logger.debug(f"Remote agent registration failed: {e}")

    def _unregister_remote_agent(self, peer_id: str):
        """
        Unregister RemoteAgent and skills when a peer disconnects.
        """
        try:
            # Unregister from orchestrator
            if self._orchestrator:
                agent_id = f"remote:{peer_id}"
                self._orchestrator.registry.unregister(agent_id)
                logger.info(f"RemoteAgent unregistered: {peer_id[:8]}...")

            # Unregister from SkillBus
            try:
                from ..bus import get_skill_bus
                from .mesh_skill_endpoint import unregister_remote_skills

                bus = get_skill_bus()
                unregister_remote_skills(bus, peer_id)
            except Exception:
                pass

        except Exception as e:
            logger.debug(f"Remote agent unregistration: {e}")


# ============================================================
# NODE - Connects to gateway
# ============================================================


class MeshNode:
    """
    A node that connects to the gateway.
    Runs on laptops, phones, other devices.

    Usage:
        node = MeshNode(
            gateway_url="ws://pi.local:18789",
            secret_key="...",
            name="macbook"
        )
        await node.connect()
    """

    def __init__(
        self, gateway_url: str, secret_key: str, name: str = "node", node_type: str = "desktop"
    ):
        self.gateway_url = gateway_url
        self.secret_key = secret_key
        self.name = name
        self.node_type = node_type
        self.node_id = secrets.token_hex(8)
        self.websocket = None
        self._running = False
        self._local_tools: dict[str, Callable] = {}
        self._message_handlers: dict[str, Callable] = {}
        self._reconnect_delay = 5

    def register_tool(self, name: str, description: str, schema: dict, handler: Callable):
        """Register a local tool that the gateway can call."""
        self._local_tools[name] = {
            "name": name,
            "description": description,
            "input_schema": schema,
            "handler": handler,
        }

    def on_message(self, msg_type: str, handler: Callable):
        """Register handler for a message type."""
        self._message_handlers[msg_type] = handler

    async def connect(self):
        """Connect to the gateway."""
        try:
            import websockets
        except ImportError:
            raise RuntimeError("websockets not available. Check server logs.")

        self._running = True

        while self._running:
            try:
                logger.info(f"Connecting to gateway: {self.gateway_url}")

                async with websockets.connect(self.gateway_url) as ws:
                    self.websocket = ws
                    self._reconnect_delay = 5  # Reset on successful connect

                    # Send hello
                    await ws.send(
                        Message(
                            type=MsgType.HELLO,
                            payload={
                                "node_id": self.node_id,
                                "name": self.name,
                                "type": self.node_type,
                                "secret_key": self.secret_key,
                                "tools": [
                                    {
                                        "name": t["name"],
                                        "description": t["description"],
                                        "input_schema": t["input_schema"],
                                    }
                                    for t in self._local_tools.values()
                                ],
                                "capabilities": ["chat", "tools"],
                            },
                        ).to_json()
                    )

                    # Wait for auth response
                    raw = await asyncio.wait_for(ws.recv(), timeout=10)
                    msg = Message.from_json(raw)

                    if msg.type == MsgType.AUTH_FAIL:
                        logger.error(f"Auth failed: {msg.payload.get('reason')}")
                        self._running = False
                        return

                    if msg.type == MsgType.AUTH_OK:
                        logger.info(f"Connected to gateway: {msg.payload.get('gateway_name')}")

                        # Start heartbeat
                        asyncio.create_task(self._heartbeat_loop())

                        # Handle messages
                        await self._message_loop(ws)

            except Exception as e:
                logger.error(f"Connection error: {e}")

                if self._running:
                    logger.info(f"Reconnecting in {self._reconnect_delay}s...")
                    await asyncio.sleep(self._reconnect_delay)
                    self._reconnect_delay = min(60, self._reconnect_delay * 2)

    async def disconnect(self):
        """Disconnect from gateway."""
        self._running = False
        if self.websocket:
            await self.websocket.close()

    async def _message_loop(self, ws):
        """Handle incoming messages."""
        async for raw in ws:
            try:
                msg = Message.from_json(raw)

                # Check for registered handler
                if msg.type in self._message_handlers:
                    await self._message_handlers[msg.type](msg)

                elif msg.type == MsgType.TOOL_REQUEST:
                    await self._handle_tool_request(ws, msg)

                elif msg.type == MsgType.ASSISTANT_RESPONSE:
                    # Default: print response
                    print(f"\nAssistant: {msg.payload.get('text', '')}\n")

            except Exception as e:
                logger.error(f"Error handling message: {e}")

    async def _handle_tool_request(self, ws, msg: Message):
        """Handle tool request from gateway."""
        tool_name = msg.payload.get("tool")
        tool_input = msg.payload.get("input", {})
        request_id = msg.payload.get("request_id")

        tool = self._local_tools.get(tool_name)

        if not tool:
            result = f"Tool not found: {tool_name}"
        else:
            try:
                handler = tool["handler"]
                if asyncio.iscoroutinefunction(handler):
                    result = await handler(tool_input)
                else:
                    result = handler(tool_input)
            except Exception as e:
                result = f"Tool error: {e}"

        await ws.send(
            Message(
                type=MsgType.TOOL_RESPONSE,
                payload={"request_id": request_id, "tool": tool_name, "result": result},
            ).to_json()
        )

    async def _heartbeat_loop(self):
        """Send periodic heartbeats."""
        while self._running and self.websocket:
            try:
                await asyncio.sleep(30)
                if self.websocket:
                    await self.websocket.send(Message(type=MsgType.HEARTBEAT, payload={}).to_json())
            except (WebSocketConnectionClosed, OSError, asyncio.CancelledError):
                break

    async def send_message(self, text: str, chat_id: str = None):
        """Send a user message to the gateway."""
        if not self.websocket:
            raise RuntimeError("Not connected")

        await self.websocket.send(
            Message(
                type=MsgType.USER_MESSAGE,
                payload={"text": text, "chat_id": chat_id or self.node_id},
            ).to_json()
        )

    async def get_nodes(self) -> list:
        """Get list of connected nodes."""
        if not self.websocket:
            return []

        await self.websocket.send(Message(type=MsgType.NODE_LIST, payload={}).to_json())

        # Wait for response (simplified)
        raw = await asyncio.wait_for(self.websocket.recv(), timeout=5)
        msg = Message.from_json(raw)

        if msg.type == MsgType.NODE_LIST:
            return msg.payload.get("nodes", [])

        return []


# ============================================================
# Helper functions
# ============================================================


def get_mesh_config() -> MeshConfig:
    """Get mesh configuration."""
    return MeshConfig()


def generate_connection_string(host: str = None, port: int = None) -> str:
    """Generate connection string for nodes."""
    config = MeshConfig()
    host = host or config.get("gateway_host", "localhost")
    port = port or config.get("gateway_port", 18789)
    key = config.get("secret_key", "")

    return f"ws://{host}:{port}?key={key[:16]}..."


async def run_gateway(agent, host: str = "0.0.0.0", port: int = 18789):
    """Convenience function to run gateway."""
    gateway = MeshGateway(agent, host, port)
    await gateway.start()

    # Keep running
    try:
        while True:
            await asyncio.sleep(1)
    except KeyboardInterrupt:
        await gateway.stop()


async def run_node(gateway_url: str, secret_key: str, name: str = "node"):
    """Convenience function to run a node."""
    node = MeshNode(gateway_url, secret_key, name)
    await node.connect()
