# 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 hashlib
import hmac
import json
import logging
import secrets
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any, Callable, Optional

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"


@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,
        }

    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._nodes_lock = asyncio.Lock()
        self._running = False
        self._server = None

        # F-01 FIX: Wire MeshTrustManager so every node connection uses
        # X3DH key exchange + Double Ratchet encryption on the wire.
        # Falls back gracefully when pynacl is not installed.
        self._trust_manager: Optional[Any] = None
        try:
            from .mesh.trust import MeshTrustManager

            self._trust_manager = MeshTrustManager(mesh_dir=MESH_DIR)
            logger.info("MeshTrustManager loaded — Signal-grade E2E encryption active")
        except Exception as e:
            logger.warning(
                f"MeshTrustManager unavailable ({e}). "
                "Messages will be transmitted without per-node E2E encryption. "
                "Install pynacl to enable Signal-protocol encryption."
            )

    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]}...")
        if self._trust_manager:
            fp = self._trust_manager.get_our_fingerprint()
            logger.info(f"Gateway identity fingerprint: {fp}")

    async def stop(self):
        """Stop the gateway."""
        self._running = False
        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
            raw = await asyncio.wait_for(websocket.recv(), timeout=10)
            msg = Message.from_json(raw)

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

            # HMAC challenge-response: the secret key is NEVER sent over the wire.
            # 1. Gateway sends a random challenge nonce.
            # 2. Node must reply with HMAC-SHA256(secret_key, challenge).
            # 3. Gateway verifies with hmac.compare_digest (constant-time).
            challenge = secrets.token_hex(32)
            await websocket.send(
                Message(type=MsgType.AUTH, payload={"challenge": challenge}).to_json()
            )

            # Wait for node's HMAC response
            raw_auth = await asyncio.wait_for(websocket.recv(), timeout=10)
            auth_msg = Message.from_json(raw_auth)
            if auth_msg.type != MsgType.AUTH:
                await websocket.close(1002, "Expected AUTH response")
                return

            expected = hmac.new(
                self.config.get("secret_key", "").encode(),
                challenge.encode(),
                hashlib.sha256,
            ).hexdigest()
            provided = auth_msg.payload.get("hmac", "")
            if not hmac.compare_digest(expected, provided):
                await websocket.send(
                    Message(
                        type=MsgType.AUTH_FAIL, payload={"reason": "Authentication failed"}
                    ).to_json()
                )
                await websocket.close(1002, "Auth failed")
                return

            # Register node
            node_id = msg.payload.get("node_id", secrets.token_hex(8))
            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),
            )

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

            # ── F-01 FIX: X3DH handshake ──────────────────────────────────
            # Correct X3DH flow:
            #   1. Node sends prekey_bundle + identity_key in HELLO
            #   2. Gateway calls initiate_handshake(node's bundle) → init_hex
            #   3. Gateway sends init_hex + its own bundle in AUTH_OK
            #   4. Node calls accept_handshake(init_hex, gw_identity_key)
            # Both sides now have a symmetric Double Ratchet session.
            auth_ok_extras = {}
            if self._trust_manager:
                node_bundle = msg.payload.get("prekey_bundle")  # aligned with HELLO key
                node_identity_key = msg.payload.get("identity_key")
                if node_bundle and node_identity_key:
                    try:
                        # Gateway initiates: consumes node's prekey bundle,
                        # returns the X3DH init message bytes to send back.
                        init_hex = self._trust_manager.initiate_handshake(
                            peer_id=node_id,
                            peer_name=node_info.name,
                            their_bundle_dict=node_bundle,
                        )
                        if init_hex:
                            auth_ok_extras["x3dh_init_message"] = init_hex
                        # Auto-approve node so session is immediately active
                        self._trust_manager.approve(node_id)
                        verification_code = self._trust_manager.get_verification_code(node_id)
                        logger.info(
                            f"X3DH initiated for {node_info.name} ({node_id[:8]}...). "
                            f"Verify with node: {verification_code}"
                        )
                    except Exception as e:
                        logger.warning(f"X3DH handshake failed for {node_id[:8]}: {e}")

                # Include our prekey bundle so the node can accept our init message.
                # BUG 3 FIX: guard _session_manager before accessing it directly.
                our_bundle = self._trust_manager.get_prekey_bundle_dict()
                if our_bundle:
                    auth_ok_extras["gateway_prekey_bundle"] = our_bundle
                    if self._trust_manager.has_crypto and self._trust_manager._session_manager:
                        auth_ok_extras["gateway_identity_key"] = (
                            self._trust_manager._session_manager.local_keys.identity_key.public.hex()
                        )
            # ── end F-01 ───────────────────────────────────────────────────

            # 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,
                        **auth_ok_extras,
                    },
                ).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:
                async with self._nodes_lock:
                    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:
                # ── F-01 FIX: decrypt inbound if an E2E session exists ─────
                if self._trust_manager and self._trust_manager.has_session(node_id):
                    try:
                        raw = self._trust_manager.decrypt_from_peer(
                            node_id, raw if isinstance(raw, bytes) else raw.encode()
                        ).decode()
                    except Exception as e:
                        logger.warning(
                            f"Decryption failed for node {node_id[:8]}: {e}. "
                            "Dropping message — possible tampering or key mismatch."
                        )
                        continue
                # ── end F-01 ──────────────────────────────────────────────

                msg = Message.from_json(raw)
                msg.node_id = node_id

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

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

                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
                    tools = msg.payload.get("tools", [])
                    async with self._nodes_lock:
                        self.node_tools[node_id] = tools
                    logger.info(f"Node {node_id} registered tools: {tools}")

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

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

    async def _send_encrypted(self, node_id: str, websocket, msg: Message) -> None:
        """
        Send a message to a node, encrypting with Double Ratchet if a
        session exists, otherwise falling back to plaintext.

        F-01: All outbound application messages go through this helper so
        the encryption path is never bypassed by accident.
        """
        wire: str = msg.to_json()
        if self._trust_manager and self._trust_manager.has_session(node_id):
            try:
                ciphertext: bytes = self._trust_manager.encrypt_for_peer(node_id, wire.encode())
                await websocket.send(ciphertext)
                return
            except Exception as e:
                # NEVER fall back to plaintext — this would enable downgrade attacks.
                logger.error(
                    f"Encryption failed for node {node_id[:8]}: {e}. "
                    "Message NOT sent — refusing plaintext fallback."
                )
                raise
        await websocket.send(wire)

    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 (encrypted if session available)
        await self._send_encrypted(
            node_id, websocket, Message(type=MsgType.TYPING, payload={"status": "thinking"})
        )

        # 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 (encrypted if session available)
        await self._send_encrypted(
            node_id,
            websocket,
            Message(
                type=MsgType.ASSISTANT_RESPONSE, payload={"text": response, "chat_id": chat_id}
            ),
        )

    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 list(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")
                async with self._nodes_lock:
                    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 list(self.connections.items()):
            if node_id not in exclude:
                try:
                    await self._send_encrypted(node_id, ws, msg)
                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."""
        ws = self.connections.get(node_id)
        if ws:
            try:
                await self._send_encrypted(node_id, ws, msg)
                return True
            except (WebSocketConnectionClosed, OSError) as e:
                logger.debug(f"Failed to send to {node_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 list(self.node_tools.items()):
            for t in tools:
                t["_source"] = f"node:{node_id}"
            all_tools.extend(tools)

        return all_tools


# ============================================================
# 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
        # F-01 FIX: node-side trust manager mirrors gateway setup
        self._trust_manager: Optional[Any] = None
        self._gateway_node_id: str = "gateway"
        try:
            from .mesh.trust import MeshTrustManager

            self._trust_manager = MeshTrustManager(mesh_dir=MESH_DIR / f"node_{self.node_id}")
        except Exception as e:
            logger.warning(f"MeshTrustManager unavailable on node ({e}). No E2E encryption.")

    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

                    # ── F-01 FIX: include our prekey bundle in HELLO ──────────
                    # BUG 1 FIX: key must be "prekey_bundle" to match gateway read.
                    hello_extras = {}
                    if self._trust_manager:
                        our_bundle = self._trust_manager.get_prekey_bundle_dict()
                        if our_bundle:
                            hello_extras["prekey_bundle"] = our_bundle  # aligned key
                            if (
                                self._trust_manager.has_crypto
                                and self._trust_manager._session_manager
                            ):
                                hello_extras["identity_key"] = (
                                    self._trust_manager._session_manager.local_keys.identity_key.public.hex()
                                )
                    # ── end F-01 ──────────────────────────────────────────

                    # Send hello (no secret key — HMAC challenge-response follows)
                    await ws.send(
                        Message(
                            type=MsgType.HELLO,
                            payload={
                                "node_id": self.node_id,
                                "name": self.name,
                                "type": self.node_type,
                                "tools": [
                                    {
                                        "name": t["name"],
                                        "description": t["description"],
                                        "input_schema": t["input_schema"],
                                    }
                                    for t in self._local_tools.values()
                                ],
                                "capabilities": ["chat", "tools"],
                                **hello_extras,
                            },
                        ).to_json()
                    )

                    # HMAC challenge-response: gateway sends a challenge,
                    # we respond with HMAC(secret_key, challenge).
                    raw_challenge = await asyncio.wait_for(ws.recv(), timeout=10)
                    challenge_msg = Message.from_json(raw_challenge)
                    if challenge_msg.type == MsgType.AUTH:
                        challenge = challenge_msg.payload.get("challenge", "")
                        response_hmac = hmac.new(
                            self.secret_key.encode(),
                            challenge.encode(),
                            hashlib.sha256,
                        ).hexdigest()
                        await ws.send(
                            Message(type=MsgType.AUTH, payload={"hmac": response_hmac}).to_json()
                        )
                    else:
                        logger.error("Unexpected message during auth handshake")
                        self._running = False
                        return

                    # Wait for auth result
                    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')}")

                        # ── F-01 FIX: complete X3DH handshake as RESPONDER ────────
                        # BUG 2 FIX: node must ACCEPT (not initiate) — gateway sent
                        # an x3dh_init_message in AUTH_OK which we now process.
                        # Correct flow: gateway initiates → node accepts.
                        if self._trust_manager:
                            self._gateway_node_id = msg.payload.get("gateway_id", "gateway")
                            x3dh_init = msg.payload.get("x3dh_init_message")
                            gw_ik = msg.payload.get("gateway_identity_key")
                            if x3dh_init and gw_ik:
                                try:
                                    self._trust_manager.accept_handshake(
                                        peer_id=self._gateway_node_id,
                                        peer_name=msg.payload.get("gateway_name", "gateway"),
                                        init_message_hex=x3dh_init,
                                        their_identity_key_hex=gw_ik,
                                    )
                                    code = self._trust_manager.get_verification_code(
                                        self._gateway_node_id
                                    )
                                    logger.info(
                                        f"X3DH session accepted from gateway. Verify: {code}"
                                    )
                                    # Auto-approve so session is immediately active
                                    self._trust_manager.approve(self._gateway_node_id)
                                    self._trust_manager.grant_permission(
                                        self._gateway_node_id, "invoke_skills"
                                    )
                                except Exception as e:
                                    logger.warning(f"Node-side X3DH accept failed: {e}")
                        # ── end F-01 ─────────────────────────────────────────────

                        # 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:
                # ── F-01 FIX: decrypt inbound from gateway ────────────────
                if self._trust_manager and self._trust_manager.has_session(self._gateway_node_id):
                    try:
                        raw = self._trust_manager.decrypt_from_peer(
                            self._gateway_node_id, raw if isinstance(raw, bytes) else raw.encode()
                        ).decode()
                    except Exception as e:
                        logger.warning(f"Node decrypt failed: {e}. Dropping message.")
                        continue
                # ── end F-01 ─────────────────────────────────────────────

                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")

        # F-01 FIX: encrypt outbound via Double Ratchet if session active
        wire: str = Message(
            type=MsgType.USER_MESSAGE, payload={"text": text, "chat_id": chat_id or self.node_id}
        ).to_json()
        if self._trust_manager and self._trust_manager.has_session(self._gateway_node_id):
            try:
                ciphertext = self._trust_manager.encrypt_for_peer(
                    self._gateway_node_id, wire.encode()
                )
                await self.websocket.send(ciphertext)
                return
            except Exception as e:
                # NEVER fall back to plaintext — refuse to send unencrypted.
                logger.error(f"Node encrypt failed: {e}. Message NOT sent.")
                raise
        await self.websocket.send(wire)

    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()
