#!/usr/bin/env python3
"""
MCP server for multi-agent coordination via shared DAG.

Nodes have: id, parents, children, status, notes, messages, data.
Multiple agents can read/write to the graph.
"""

import asyncio
import json
import os
from datetime import datetime
from pathlib import Path

from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent


def get_graph_path():
    """Get path to graph.json file."""
    path = os.environ.get("BLACKBOARD_GRAPH_PATH")
    if path:
        return Path(path)
    return Path.cwd() / "graph.json"


def load_graph():
    """Load graph from JSON file."""
    path = get_graph_path()
    if path.exists():
        return json.loads(path.read_text())
    return {
        "nodes": {
            "root": {
                "id": "root",
                "parents": [],
                "children": [],
                "status": "active",
                "notes": [],
                "messages": [],
                "data": {}
            }
        }
    }


def save_graph(graph):
    """Save graph to JSON file."""
    path = get_graph_path()
    path.write_text(json.dumps(graph, indent=2) + "\n")


def get_node(graph, node_id):
    """Get a node by id."""
    return graph["nodes"].get(node_id)


def ensure_node_fields(node):
    """Ensure node has all required fields."""
    node.setdefault("parents", [])
    node.setdefault("children", [])
    node.setdefault("status", "active")
    node.setdefault("notes", [])
    node.setdefault("messages", [])
    node.setdefault("data", {})


def format_node(node, show_notes=True, show_messages=True, show_data=True):
    """Format a single node for display."""
    parts = [f"{node['id']} [{node.get('status', 'active')}]"]
    if node.get("parents"):
        parts.append(f"  parents: {', '.join(node['parents'])}")
    if node.get("children"):
        parts.append(f"  children: {', '.join(node['children'])}")
    if show_notes and node.get("notes"):
        parts.append(f"  notes: {len(node['notes'])}")
        for n in node["notes"]:
            parts.append(f"    - {n[:50]}{'...' if len(n) > 50 else ''}")
    if show_messages and node.get("messages"):
        parts.append(f"  messages: {len(node['messages'])}")
        for m in node["messages"]:
            parts.append(f"    - [{m.get('from', '?')}] {m['text'][:40]}{'...' if len(m['text']) > 40 else ''}")
    if show_data and node.get("data"):
        parts.append(f"  data: {json.dumps(node['data'])[:60]}{'...' if len(json.dumps(node['data'])) > 60 else ''}")
    return "\n".join(parts)


def format_graph(graph, node_id=None, show_notes=True, show_messages=True, show_data=True):
    """Format graph or single node for display."""
    if node_id:
        node = get_node(graph, node_id)
        if not node:
            return f"Error: node '{node_id}' not found"
        return format_node(node, show_notes, show_messages, show_data)

    lines = []
    for nid, node in graph["nodes"].items():
        lines.append(format_node(node, show_notes, show_messages, show_data))
        lines.append("")
    return "\n".join(lines).strip()


async def run_server():
    """Run the MCP server."""
    server = Server("mcp-blackboard-graph")

    tools = [
        Tool(
            name="show",
            description="Show the graph or a single node",
            inputSchema={
                "type": "object",
                "properties": {
                    "node": {"type": "string", "description": "Node id to show (omit for all)"},
                    "hide_notes": {"type": "boolean", "description": "Hide notes", "default": False},
                    "hide_messages": {"type": "boolean", "description": "Hide messages", "default": False},
                    "hide_data": {"type": "boolean", "description": "Hide data", "default": False}
                },
                "required": []
            }
        ),
        Tool(
            name="add",
            description="Add a new node to the graph",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "Node id (unique)"},
                    "parents": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "Parent node ids (for DAG structure)"
                    },
                    "status": {"type": "string", "description": "Initial status (default: active)"},
                    "data": {"type": "object", "description": "Initial data dict"}
                },
                "required": ["id"]
            }
        ),
        Tool(
            name="remove",
            description="Remove a node from the graph",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "Node id to remove"}
                },
                "required": ["id"]
            }
        ),
        Tool(
            name="get",
            description="Get full details of a node as JSON",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "Node id"}
                },
                "required": ["id"]
            }
        ),
        Tool(
            name="set_status",
            description="Set status of a node",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "Node id"},
                    "status": {"type": "string", "description": "New status"}
                },
                "required": ["id", "status"]
            }
        ),
        Tool(
            name="add_data",
            description="Merge data dict into a node's data",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "Node id"},
                    "data": {"type": "object", "description": "Data to merge"}
                },
                "required": ["id", "data"]
            }
        ),
        Tool(
            name="add_note",
            description="Add a note to a node",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "Node id"},
                    "note": {"type": "string", "description": "Note text"}
                },
                "required": ["id", "note"]
            }
        ),
        Tool(
            name="send_message",
            description="Send a message to a node",
            inputSchema={
                "type": "object",
                "properties": {
                    "id": {"type": "string", "description": "Node id"},
                    "message": {"type": "string", "description": "Message text"},
                    "from_agent": {"type": "string", "description": "Sender name (optional)"}
                },
                "required": ["id", "message"]
            }
        ),
    ]

    @server.list_tools()
    async def list_tools():
        return tools

    @server.call_tool()
    async def call_tool(name: str, arguments: dict):
        graph = load_graph()

        if name == "show":
            node_id = arguments.get("node")
            show_notes = not arguments.get("hide_notes", False)
            show_messages = not arguments.get("hide_messages", False)
            show_data = not arguments.get("hide_data", False)
            result = format_graph(graph, node_id, show_notes, show_messages, show_data)
            return [TextContent(type="text", text=result)]

        elif name == "add":
            node_id = arguments.get("id", "").strip()
            parents = arguments.get("parents", [])
            status = arguments.get("status", "active")
            data = arguments.get("data", {})

            if not node_id:
                return [TextContent(type="text", text="Error: id is required")]
            if node_id in graph["nodes"]:
                return [TextContent(type="text", text=f"Error: node '{node_id}' already exists")]

            # Validate parents exist
            for p in parents:
                if p not in graph["nodes"]:
                    return [TextContent(type="text", text=f"Error: parent '{p}' not found")]

            # Create node
            graph["nodes"][node_id] = {
                "id": node_id,
                "parents": parents,
                "children": [],
                "status": status,
                "notes": [],
                "messages": [],
                "data": data
            }

            # Update parent nodes' children lists
            for p in parents:
                if node_id not in graph["nodes"][p]["children"]:
                    graph["nodes"][p]["children"].append(node_id)

            save_graph(graph)
            return [TextContent(type="text", text=f"Added node '{node_id}'")]

        elif name == "remove":
            node_id = arguments.get("id", "")
            if node_id == "root":
                return [TextContent(type="text", text="Error: cannot remove root")]
            if node_id not in graph["nodes"]:
                return [TextContent(type="text", text=f"Error: node '{node_id}' not found")]

            node = graph["nodes"][node_id]

            # Remove from parents' children lists
            for p in node.get("parents", []):
                if p in graph["nodes"]:
                    graph["nodes"][p]["children"] = [c for c in graph["nodes"][p]["children"] if c != node_id]

            # Remove from children's parents lists
            for c in node.get("children", []):
                if c in graph["nodes"]:
                    graph["nodes"][c]["parents"] = [p for p in graph["nodes"][c]["parents"] if p != node_id]

            del graph["nodes"][node_id]
            save_graph(graph)
            return [TextContent(type="text", text=f"Removed node '{node_id}'")]

        elif name == "get":
            node_id = arguments.get("id", "")
            node = get_node(graph, node_id)
            if not node:
                return [TextContent(type="text", text=f"Error: node '{node_id}' not found")]
            return [TextContent(type="text", text=json.dumps(node, indent=2))]

        elif name == "set_status":
            node_id = arguments.get("id", "")
            status = arguments.get("status", "")
            node = get_node(graph, node_id)
            if not node:
                return [TextContent(type="text", text=f"Error: node '{node_id}' not found")]
            node["status"] = status
            save_graph(graph)
            return [TextContent(type="text", text=f"Set status of '{node_id}' to '{status}'")]

        elif name == "add_data":
            node_id = arguments.get("id", "")
            data = arguments.get("data", {})
            node = get_node(graph, node_id)
            if not node:
                return [TextContent(type="text", text=f"Error: node '{node_id}' not found")]
            ensure_node_fields(node)
            node["data"].update(data)
            save_graph(graph)
            return [TextContent(type="text", text=f"Added data to '{node_id}'")]

        elif name == "add_note":
            node_id = arguments.get("id", "")
            note = arguments.get("note", "")
            node = get_node(graph, node_id)
            if not node:
                return [TextContent(type="text", text=f"Error: node '{node_id}' not found")]
            ensure_node_fields(node)
            node["notes"].append(note)
            save_graph(graph)
            return [TextContent(type="text", text=f"Added note to '{node_id}'")]

        elif name == "send_message":
            node_id = arguments.get("id", "")
            message = arguments.get("message", "")
            from_agent = arguments.get("from_agent", "anonymous")
            node = get_node(graph, node_id)
            if not node:
                return [TextContent(type="text", text=f"Error: node '{node_id}' not found")]
            ensure_node_fields(node)
            node["messages"].append({
                "from": from_agent,
                "text": message,
                "time": datetime.now().isoformat()
            })
            save_graph(graph)
            return [TextContent(type="text", text=f"Message sent to '{node_id}'")]

        return [TextContent(type="text", text=f"Unknown tool: {name}")]

    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


def main():
    asyncio.run(run_server())


if __name__ == "__main__":
    main()
