# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Licensed under the MIT License

"""
Mesh Memory Bridge — Phase 3 (v2.7.2)

Enables opt-in memory sharing across trusted mesh peers.
Rides on the existing SkillBus (core/bus.py) pub/sub system and
the existing Memory class (core/memory.py) semantic search.

Design principles:
  - Memories are PRIVATE by default. User must explicitly share.
  - Queries sent as TEXT, not vectors (embedding compatibility).
  - Each peer runs their own embedding provider for local search.
  - Results tagged with origin node for transparency.
  - 2-second timeout on mesh searches (never block conversation).
  - HIPAA tenants: memory sharing disabled entirely.

Event topics (only mesh.* topics forwarded across gateways):
  mesh.memory.query         — incoming query from a peer
  mesh.memory.query_response — results sent back to querying peer
  mesh.memory.push          — peer pushing a shared memory to us
  mesh.memory.revoke        — peer revoking a previously shared memory

Integration:
  Called from agent.py when building context. Agent calls
  bridge.search_mesh() which fans out to trusted peers and
  merges results with local context.
"""

from __future__ import annotations

import asyncio
import hashlib
import logging
import threading
import time
from concurrent.futures import ThreadPoolExecutor
from dataclasses import asdict, dataclass
from datetime import datetime, timezone
from pathlib import Path
from typing import TYPE_CHECKING, Any, Dict, List, Optional

if TYPE_CHECKING:
    pass

logger = logging.getLogger(__name__)

MESH_MEMORY_QUERY = "mesh.memory.query"
MESH_MEMORY_RESPONSE = "mesh.memory.query_response"
MESH_MEMORY_PUSH = "mesh.memory.push"
MESH_MEMORY_REVOKE = "mesh.memory.revoke"

# Default timeout for mesh memory searches (seconds)
MESH_SEARCH_TIMEOUT = 2.0
# Maximum results per peer
MAX_RESULTS_PER_PEER = 5


def _utcnow() -> datetime:
    return datetime.now(timezone.utc)


# ============================================================
# SHARED MEMORY METADATA
# ============================================================


@dataclass
class SharedMemoryMeta:
    """
    Metadata tracking which memories are shared and with whom.
    Stored alongside normal memories in a separate sidecar file.
    """

    key: str
    shareable: bool = False
    share_scope: str = "all"  # "all" or comma-separated node_ids
    mesh_origin: Optional[str] = None  # node_id if received from a peer
    mesh_origin_name: Optional[str] = None  # display name of origin node
    mesh_ttl: int = 3600  # seconds before cache expires
    shared_at: Optional[str] = None
    received_at: Optional[str] = None

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "SharedMemoryMeta":
        known = {f.name for f in cls.__dataclass_fields__.values()}
        return cls(**{k: v for k, v in d.items() if k in known})

    def is_shared_with(self, node_id: str) -> bool:
        """Check if this memory is shared with a specific node."""
        if not self.shareable:
            return False
        if self.share_scope == "all":
            return True
        return node_id in self.share_scope.split(",")


# ============================================================
# SHARED MEMORY STORE — sidecar for Memory
# ============================================================

import json  # noqa: E402


class SharedMemoryStore:
    """
    Sidecar store tracking which memories are shared.
    Stored in MESH_DIR/shared_memories.json.
    Does NOT store the memories themselves — those stay in Memory.
    """

    def __init__(self, storage_path: Path):
        self._path = storage_path / "shared_memories.json"
        self._lock = threading.Lock()
        self._meta: Dict[str, SharedMemoryMeta] = {}
        self._load()

    def _load(self):
        if self._path.exists():
            try:
                data = json.loads(self._path.read_text())
                for entry in data.get("shared", []):
                    meta = SharedMemoryMeta.from_dict(entry)
                    self._meta[meta.key] = meta
            except (json.JSONDecodeError, IOError, OSError) as e:
                logger.warning(f"Failed to load shared memory store: {e}")

    def _save(self):
        try:
            self._path.parent.mkdir(parents=True, exist_ok=True)
            data = {"shared": [m.to_dict() for m in self._meta.values()]}
            self._path.write_text(json.dumps(data, indent=2))
        except (IOError, OSError) as e:
            logger.warning(f"Failed to save shared memory store: {e}")

    def share(self, key: str, scope: str = "all") -> SharedMemoryMeta:
        """Mark a memory as shareable."""
        with self._lock:
            meta = self._meta.get(key, SharedMemoryMeta(key=key))
            meta.shareable = True
            meta.share_scope = scope
            meta.shared_at = _utcnow().isoformat()
            self._meta[key] = meta
            self._save()
            return meta

    def unshare(self, key: str) -> bool:
        """Remove sharing from a memory."""
        with self._lock:
            if key in self._meta:
                self._meta[key].shareable = False
                self._save()
                return True
            return False

    def store_received(
        self, key: str, origin_id: str, origin_name: str, ttl: int = 3600
    ) -> SharedMemoryMeta:
        """Track a memory received from a peer."""
        with self._lock:
            meta = SharedMemoryMeta(
                key=key,
                shareable=False,  # Received memories not re-shared by default
                mesh_origin=origin_id,
                mesh_origin_name=origin_name,
                mesh_ttl=ttl,
                received_at=_utcnow().isoformat(),
            )
            self._meta[key] = meta
            self._save()
            return meta

    def get(self, key: str) -> Optional[SharedMemoryMeta]:
        with self._lock:
            return self._meta.get(key)

    def get_shareable_keys(self, for_node: Optional[str] = None) -> List[str]:
        """Get keys of all shareable memories, optionally filtered by node."""
        with self._lock:
            keys = []
            for key, meta in self._meta.items():
                if meta.shareable:
                    if for_node is None or meta.is_shared_with(for_node):
                        keys.append(key)
            return keys

    def get_received_keys(self) -> List[str]:
        """Get keys of memories received from peers."""
        with self._lock:
            return [k for k, m in self._meta.items() if m.mesh_origin]

    def remove(self, key: str) -> bool:
        with self._lock:
            if key in self._meta:
                del self._meta[key]
                self._save()
                return True
            return False

    def count_shared(self) -> int:
        with self._lock:
            return len([m for m in self._meta.values() if m.shareable])

    def count_received(self) -> int:
        with self._lock:
            return len([m for m in self._meta.values() if m.mesh_origin])


# ============================================================
# MESH SEARCH RESULT
# ============================================================


@dataclass
class MeshMemoryResult:
    """A single memory result from a mesh peer."""

    key: str
    value: str
    category: str
    importance: int
    origin_node_id: str
    origin_node_name: str
    relevance_score: float = 0.0

    @property
    def tagged_value(self) -> str:
        """Value with origin tag for context injection."""
        return f"[{self.origin_node_name}] {self.value}"

    def to_dict(self) -> dict:
        return asdict(self)

    @classmethod
    def from_dict(cls, d: dict) -> "MeshMemoryResult":
        known = {f.name for f in cls.__dataclass_fields__.values()}
        return cls(**{k: v for k, v in d.items() if k in known})


# ============================================================
# MEMORY BRIDGE — the main class
# ============================================================


class MemoryBridge:
    """
    Bridges local Memory with mesh peers for shared context.

    Handles:
      - Publishing shareable memories in response to peer queries
      - Querying peers for context relevant to the current conversation
      - Receiving pushed memories from peers
      - Revoking shared memories

    Usage:
        bridge = MemoryBridge(
            memory=agent.memory,
            node_id="local_123",
            node_name="Kitchen Pi",
            mesh_dir=MESH_DIR,
        )

        # Wire into bus
        bridge.register_bus_handlers(skill_bus)

        # In agent context building:
        mesh_results = bridge.search_mesh(
            "What dietary restrictions?",
            peer_ids=["peer_abc"],
            timeout=2.0,
        )
    """

    def __init__(
        self,
        memory: Any,  # Memory instance
        node_id: str,
        node_name: str,
        mesh_dir: Optional[Path] = None,
        trust_manager: Any = None,  # MeshTrustManager
        hipaa_enabled: bool = False,
    ):
        self.memory = memory
        self.node_id = node_id
        self.node_name = node_name
        self.trust_manager = trust_manager
        self.hipaa_enabled = hipaa_enabled

        if mesh_dir is None:
            mesh_dir = Path.home() / ".familiar" / "data" / "mesh"
        self.shared_store = SharedMemoryStore(mesh_dir)

        # Pending query responses (query_id -> list of results)
        self._pending_responses: Dict[str, List[MeshMemoryResult]] = {}
        self._response_events: Dict[str, threading.Event] = {}
        self._lock = threading.Lock()

        # Gateway reference (set by MeshGateway after init)
        self._gateway = None

        # Thread pool for parallel peer queries
        self._executor = ThreadPoolExecutor(max_workers=4, thread_name_prefix="mesh_mem")

        if hipaa_enabled:
            logger.info("HIPAA mode: memory sharing disabled")

    def set_gateway(self, gateway):
        """Set reference to MeshGateway for sending messages."""
        self._gateway = gateway

    # ── Sharing commands ──

    def share_memory(self, key: str, scope: str = "all") -> bool:
        """
        Mark a local memory as shareable with mesh peers.

        Args:
            key: Memory key to share
            scope: "all" for all trusted peers, or comma-separated node_ids
        """
        if self.hipaa_enabled:
            logger.warning("Cannot share memories in HIPAA mode")
            return False

        if key not in self.memory.memories:
            logger.warning(f"Memory key '{key}' not found")
            return False

        self.shared_store.share(key, scope)
        logger.info(f"Memory shared: '{key}' (scope={scope})")
        return True

    def unshare_memory(self, key: str) -> bool:
        """Remove sharing from a memory and notify peers."""
        result = self.shared_store.unshare(key)
        if result:
            logger.info(f"Memory unshared: '{key}'")
            # Broadcast revocation to peers
            self._broadcast_revocation(key)
        return result

    def get_shared_keys(self) -> List[str]:
        """Get all shared memory keys."""
        return self.shared_store.get_shareable_keys()

    # ── Query handling (responding to peer queries) ──

    def handle_remote_query(self, message: dict) -> Optional[dict]:
        """
        Handle an incoming memory query from a peer.
        Called when a mesh.memory.query event arrives.

        Args:
            message: {query_text, requesting_node, query_id, max_results}

        Returns:
            Response dict with results, or None if denied
        """
        if self.hipaa_enabled:
            return None

        requesting_node = message.get("requesting_node", "")

        # Check permission
        if self.trust_manager and not self.trust_manager.check_permission(
            requesting_node, "request_memory"
        ):
            logger.debug(f"Memory query denied for {requesting_node}: no request_memory permission")
            return None

        query_text = message.get("query_text", "")
        max_results = min(message.get("max_results", MAX_RESULTS_PER_PEER), MAX_RESULTS_PER_PEER)
        query_id = message.get("query_id", "")

        # Search only shareable memories
        shareable_keys = set(self.shared_store.get_shareable_keys(for_node=requesting_node))
        if not shareable_keys:
            return {
                "query_id": query_id,
                "results": [],
                "origin_node": self.node_id,
                "origin_name": self.node_name,
            }

        # Run local search
        all_results = self.memory.search(query_text, max_results=max_results * 2)

        # Filter to only shareable
        filtered = [r for r in all_results if r.key in shareable_keys][:max_results]

        results = [
            {
                "key": r.key,
                "value": r.value,
                "category": r.category,
                "importance": r.importance,
            }
            for r in filtered
        ]

        return {
            "query_id": query_id,
            "results": results,
            "origin_node": self.node_id,
            "origin_name": self.node_name,
        }

    def handle_remote_response(self, message: dict):
        """
        Handle a query response from a peer.
        Called when a mesh.memory.query_response arrives.
        """
        query_id = message.get("query_id", "")
        origin_node = message.get("origin_node", "")
        origin_name = message.get("origin_name", origin_node)

        results = []
        for r in message.get("results", []):
            results.append(
                MeshMemoryResult(
                    key=r.get("key", ""),
                    value=r.get("value", ""),
                    category=r.get("category", "fact"),
                    importance=r.get("importance", 5),
                    origin_node_id=origin_node,
                    origin_node_name=origin_name,
                )
            )

        with self._lock:
            if query_id in self._pending_responses:
                self._pending_responses[query_id].extend(results)
                event = self._response_events.get(query_id)
                if event:
                    event.set()

    def handle_remote_push(self, message: dict):
        """Handle a peer pushing a shared memory to us."""
        if self.hipaa_enabled:
            return

        origin_node = message.get("origin_node", "")

        # Check permission
        if self.trust_manager and not self.trust_manager.check_permission(
            origin_node, "request_memory"
        ):
            return

        key = message.get("key", "")
        value = message.get("value", "")
        category = message.get("category", "fact")
        origin_name = message.get("origin_name", origin_node)
        ttl = message.get("ttl", 3600)

        if not key or not value:
            return

        # Store as a received memory with mesh origin tag
        mesh_key = f"mesh:{origin_node[:8]}:{key}"
        self.memory.remember(
            mesh_key,
            value,
            category=category,
            source=f"mesh:{origin_name}",
            importance=3,  # Mesh memories start at lower importance
        )
        self.shared_store.store_received(mesh_key, origin_node, origin_name, ttl)
        logger.debug(f"Received memory from {origin_name}: {key}")

    def handle_remote_revoke(self, message: dict):
        """Handle a peer revoking a previously shared memory."""
        origin_node = message.get("origin_node", "")
        key = message.get("key", "")
        mesh_key = f"mesh:{origin_node[:8]}:{key}"

        if mesh_key in self.memory.memories:
            self.memory.forget(mesh_key)
            self.shared_store.remove(mesh_key)
            logger.debug(f"Revoked mesh memory: {mesh_key}")

    # ── Outgoing queries ──

    def search_mesh(
        self,
        query: str,
        peer_ids: Optional[List[str]] = None,
        max_results: int = MAX_RESULTS_PER_PEER,
        timeout: float = MESH_SEARCH_TIMEOUT,
    ) -> List[MeshMemoryResult]:
        """
        Search memories across mesh peers.

        Sends query to all trusted peers (or specific peer_ids),
        waits up to timeout seconds, merges results.

        Args:
            query: Natural language query
            peer_ids: Specific peers to query (None = all trusted)
            max_results: Max total results
            timeout: Seconds to wait (default 2.0)

        Returns:
            List of MeshMemoryResult sorted by importance
        """
        if self.hipaa_enabled:
            return []

        if not self._gateway:
            return []

        # Determine target peers
        targets = peer_ids or self._get_queryable_peers()
        if not targets:
            return []

        # Create query with unique ID
        query_id = hashlib.sha256(f"{query}:{time.time()}:{self.node_id}".encode()).hexdigest()[:16]

        # Set up response collection
        with self._lock:
            self._pending_responses[query_id] = []
            self._response_events[query_id] = threading.Event()

        # Send query to each peer
        query_msg = {
            "query_text": query,
            "requesting_node": self.node_id,
            "query_id": query_id,
            "max_results": max_results,
        }

        for peer_id in targets:
            try:
                self._send_to_peer(peer_id, MESH_MEMORY_QUERY, query_msg)
            except Exception as e:
                logger.debug(f"Failed to send query to {peer_id}: {e}")

        # Wait for responses (up to timeout)
        event = self._response_events[query_id]
        event.wait(timeout=timeout)

        # Small grace period for additional responses
        time.sleep(min(0.1, timeout * 0.05))

        # Collect and clean up
        with self._lock:
            results = self._pending_responses.pop(query_id, [])
            self._response_events.pop(query_id, None)

        # Sort by importance and deduplicate
        seen_values = set()
        deduped = []
        for r in sorted(results, key=lambda x: x.importance, reverse=True):
            if r.value not in seen_values:
                seen_values.add(r.value)
                deduped.append(r)

        return deduped[:max_results]

    def get_mesh_context(
        self,
        query: str,
        timeout: float = MESH_SEARCH_TIMEOUT,
    ) -> str:
        """
        Get mesh memory context formatted for LLM prompt injection.

        Returns a string like:
            ## Mesh Context
            - [Kitchen Pi] User is gluten-free
            - [Office Desktop] User prefers dark roast coffee
        """
        results = self.search_mesh(query, timeout=timeout)
        if not results:
            return ""

        lines = ["<mesh_context>"]
        for r in results:
            lines.append(f"- {r.tagged_value}")
        lines.append("</mesh_context>")

        return "\n".join(lines)

    # ── Bus registration ──

    def register_bus_handlers(self, bus):
        """Register event handlers on the SkillBus."""
        bus.subscribe(MESH_MEMORY_QUERY, self._bus_handle_query, skill="mesh_memory")
        bus.subscribe(MESH_MEMORY_RESPONSE, self._bus_handle_response, skill="mesh_memory")
        bus.subscribe(MESH_MEMORY_PUSH, self._bus_handle_push, skill="mesh_memory")
        bus.subscribe(MESH_MEMORY_REVOKE, self._bus_handle_revoke, skill="mesh_memory")
        logger.info("Memory bridge: bus handlers registered")

    def _bus_handle_query(self, message):
        """Bus handler for incoming queries."""
        response = self.handle_remote_query(message.payload)
        if response and self._gateway:
            requesting_node = message.payload.get("requesting_node", "")
            self._send_to_peer(requesting_node, MESH_MEMORY_RESPONSE, response)

    def _bus_handle_response(self, message):
        """Bus handler for incoming responses."""
        self.handle_remote_response(message.payload)

    def _bus_handle_push(self, message):
        """Bus handler for pushed memories."""
        self.handle_remote_push(message.payload)

    def _bus_handle_revoke(self, message):
        """Bus handler for revocations."""
        self.handle_remote_revoke(message.payload)

    # ── Internal helpers ──

    def _get_queryable_peers(self) -> List[str]:
        """Get peer IDs that have request_memory permission."""
        if not self.trust_manager:
            # Without trust manager, query all connected peers
            if self._gateway:
                return list(self._gateway.peer_gateways.keys())
            return []

        trusted = self.trust_manager.get_trusted_peers()
        return [r.node_id for r in trusted if r.permissions.check("request_memory")]

    def _send_to_peer(self, peer_id: str, topic: str, payload: dict):
        """Send a mesh event to a specific peer via the gateway."""
        if not self._gateway:
            return

        from .gateway import Message, MsgType

        msg = Message(
            type=MsgType.TOOL_REQUEST,
            payload={
                "mesh_event": True,
                "topic": topic,
                "data": payload,
            },
        )

        ws = self._gateway.peer_gateways.get(peer_id)
        if ws:
            # We're in a sync context but ws.send is async
            # Use the gateway's event loop
            try:
                loop = asyncio.get_event_loop()
                if loop.is_running():
                    asyncio.ensure_future(ws.send(msg.to_json()))
                else:
                    loop.run_until_complete(ws.send(msg.to_json()))
            except RuntimeError:
                # No event loop — create one
                asyncio.run(ws.send(msg.to_json()))

    def _broadcast_revocation(self, key: str):
        """Broadcast a memory revocation to all connected peers."""
        if not self._gateway:
            return

        revoke_msg = {
            "origin_node": self.node_id,
            "key": key,
        }
        for peer_id in self._gateway.peer_gateways:
            try:
                self._send_to_peer(peer_id, MESH_MEMORY_REVOKE, revoke_msg)
            except Exception as e:
                logger.debug(f"Failed to send revocation to {peer_id}: {e}")

    # ── Status ──

    def get_status(self) -> dict:
        """Get memory bridge status for CLI/dashboard."""
        return {
            "hipaa_mode": self.hipaa_enabled,
            "shared_count": self.shared_store.count_shared(),
            "received_count": self.shared_store.count_received(),
            "shared_keys": self.get_shared_keys(),
            "queryable_peers": len(self._get_queryable_peers()),
            "gateway_connected": self._gateway is not None,
        }
