"""LangChain tools for Mesh Cognition shared memory."""

from typing import Optional, Type

from langchain_core.callbacks import CallbackManagerForToolRun
from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

from .wrapper import MeshCognitionAPIWrapper


# ── Input Schemas ───────────────────────────────────────────────────

class MeshWriteInput(BaseModel):
    content: str = Field(description="The content to remember and share across the mesh")
    key: str = Field(default="", description="Optional key for the memory entry")
    tags: str = Field(default="", description="Comma-separated tags for categorization")


class MeshSearchInput(BaseModel):
    query: str = Field(description="Search query to find memories across the mesh")


class MeshStatusInput(BaseModel):
    pass


class MeshContextInput(BaseModel):
    pass


# ── Tools ───────────────────────────────────────────────────────────

class MeshMemoryWrite(BaseTool):
    """Write a memory to the mesh. All connected peers can recall it."""

    name: str = "mesh_memory_write"
    description: str = (
        "Store a memory in the mesh network. The memory is saved locally and "
        "broadcast to all connected peers. Use this when you learn something "
        "that other agents on the mesh should know — facts, decisions, plans, "
        "meeting notes, or any knowledge worth sharing."
    )
    args_schema: Type[BaseModel] = MeshWriteInput
    api: MeshCognitionAPIWrapper = Field(default_factory=MeshCognitionAPIWrapper)

    def _run(
        self,
        content: str,
        key: str = "",
        tags: str = "",
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        tag_list = [t.strip() for t in tags.split(",") if t.strip()] if tags else []
        result = self.api.write_memory(content, key=key or None, tags=tag_list or None)
        peers = result.get("broadcasted", 0)
        return f"Stored and broadcast to {peers} peer(s)."


class MeshMemorySearch(BaseTool):
    """Search memories across the entire mesh — local and all peers."""

    name: str = "mesh_memory_search"
    description: str = (
        "Search for memories across the mesh network. Returns results from "
        "this device AND all connected peers. Use this when you need to recall "
        "something — it may have been stored by another agent on a different device."
    )
    args_schema: Type[BaseModel] = MeshSearchInput
    api: MeshCognitionAPIWrapper = Field(default_factory=MeshCognitionAPIWrapper)

    def _run(
        self,
        query: str,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        results = self.api.search_memory(query)
        if not results:
            return "No memories found."
        lines = []
        for r in results:
            source = r.get("_source", r.get("source", "unknown"))
            content = r.get("content", "")
            tags = ", ".join(r.get("tags", []))
            lines.append(f"[{source}] {content}" + (f" (tags: {tags})" if tags else ""))
        return "\n".join(lines)


class MeshStatus(BaseTool):
    """Check the mesh network status — peers, coupling, synchronization."""

    name: str = "mesh_status"
    description: str = (
        "Get the current mesh network status including active peers, "
        "Kuramoto synchronization parameter r(t), and coupling decisions. "
        "Use this to understand the health and state of the mesh."
    )
    args_schema: Type[BaseModel] = MeshStatusInput
    api: MeshCognitionAPIWrapper = Field(default_factory=MeshCognitionAPIWrapper)

    def _run(
        self,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        status = self.api.status()
        peers = status.get("activePeers", 0)
        r = status.get("kuramotoR")
        hostname = status.get("hostname", "unknown")
        uptime = int(status.get("uptime", 0))
        r_str = f"{r:.3f}" if r is not None else "N/A"
        return (
            f"Mesh node: {hostname}\n"
            f"Active peers: {peers}\n"
            f"Kuramoto r(t): {r_str}\n"
            f"Uptime: {uptime}s"
        )


class MeshContext(BaseTool):
    """Get mesh-influenced context from aligned peers."""

    name: str = "mesh_context"
    description: str = (
        "Get context from aligned mesh peers — what other agents on the "
        "network are focused on. Use this to enrich your reasoning with "
        "collective intelligence from the mesh."
    )
    args_schema: Type[BaseModel] = MeshContextInput
    api: MeshCognitionAPIWrapper = Field(default_factory=MeshCognitionAPIWrapper)

    def _run(
        self,
        run_manager: Optional[CallbackManagerForToolRun] = None,
    ) -> str:
        ctx = self.api.context()
        lines = [f"Local context: {ctx.get('localContext', 'None')}"]
        for peer in ctx.get("alignedPeers", []):
            lines.append(f"Peer {peer.get('hostname', peer.get('peerId', '?'))}: {peer.get('decision', '?')}")
        for mem in ctx.get("peerMemories", []):
            lines.append(f"Peer memory [{mem.get('source', '?')}]: {mem.get('content', '')}")
        return "\n".join(lines)
