"""
Phase 59: Knowledge Graph MCP Plugin.

Provides a shared memory graph for the swarm. Allows agents to create entities,
link them via relations, and query the graph to understand system architecture,
dependencies, and contextual knowledge across the decentralized network.
"""

import json
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List

from mcp.types import Tool, TextContent
from brynq.server import BROKER_DIR, _write_json_sync, SESSION_ID

# ── Paths ──────────────────────────────────────────────────────────────────
KNOWLEDGE_DIR = BROKER_DIR / "knowledge_graph"
GRAPH_FILE = KNOWLEDGE_DIR / "graph.json"

KNOWLEDGE_DIR.mkdir(parents=True, exist_ok=True)

# ── Storage Helpers ────────────────────────────────────────────────────────
def _read_graph() -> dict:
    if not GRAPH_FILE.exists():
        return {"entities": {}, "relations": []}
    try:
        return json.loads(GRAPH_FILE.read_text("utf-8"))
    except (json.JSONDecodeError, OSError):
        return {"entities": {}, "relations": []}

def _save_graph(graph: dict) -> None:
    _write_json_sync(GRAPH_FILE, graph)

# ── Tool Handlers ──────────────────────────────────────────────────────────

async def _handle_create_entity(args: dict) -> list[TextContent]:
    entity_id = args.get("id")
    entity_type = args.get("type", "concept")
    attributes = args.get("attributes", {})
    
    if not entity_id:
        return [TextContent(type="text", text="Error: 'id' is required")]
        
    graph = _read_graph()
    
    graph["entities"][entity_id] = {
        "id": entity_id,
        "type": entity_type,
        "attributes": attributes,
        "created_by": SESSION_ID,
        "created_at": datetime.now(timezone.utc).isoformat(),
    }
    
    _save_graph(graph)
    return [TextContent(type="text", text=f"Entity '{entity_id}' created/updated successfully.")]

async def _handle_create_relation(args: dict) -> list[TextContent]:
    source = args.get("source")
    target = args.get("target")
    relation_type = args.get("relation_type", "relates_to")
    
    if not source or not target:
        return [TextContent(type="text", text="Error: 'source' and 'target' are required")]
        
    graph = _read_graph()
    
    if source not in graph["entities"]:
        return [TextContent(type="text", text=f"Error: Source entity '{source}' not found")]
    if target not in graph["entities"]:
        return [TextContent(type="text", text=f"Error: Target entity '{target}' not found")]
        
    # Prevent exact duplicates
    for rel in graph["relations"]:
        if rel["source"] == source and rel["target"] == target and rel["type"] == relation_type:
            return [TextContent(type="text", text="Relation already exists.")]
            
    relation = {
        "source": source,
        "target": target,
        "type": relation_type,
        "created_by": SESSION_ID,
        "created_at": datetime.now(timezone.utc).isoformat(),
    }
    
    graph["relations"].append(relation)
    _save_graph(graph)
    return [TextContent(type="text", text=f"Relation '{source}' -[{relation_type}]-> '{target}' created.")]

async def _handle_query_graph(args: dict) -> list[TextContent]:
    query_id = args.get("entity_id")
    
    graph = _read_graph()
    
    if not query_id:
        # Return summary if no specific entity queried
        summary = {
            "entity_count": len(graph["entities"]),
            "relation_count": len(graph["relations"]),
            "entities": list(graph["entities"].keys())
        }
        return [TextContent(type="text", text=json.dumps(summary, indent=2))]
        
    if query_id not in graph["entities"]:
        return [TextContent(type="text", text=f"Entity '{query_id}' not found.")]
        
    entity = graph["entities"][query_id]
    
    # Find incoming and outgoing relations
    outgoing = [r for r in graph["relations"] if r["source"] == query_id]
    incoming = [r for r in graph["relations"] if r["target"] == query_id]
    
    result = {
        "entity": entity,
        "relations": {
            "outgoing": outgoing,
            "incoming": incoming
        }
    }
    
    return [TextContent(type="text", text=json.dumps(result, indent=2))]


# ── MCP Plugin Exports ─────────────────────────────────────────────────────

TOOLS = [
    Tool(
        name="kg_create_entity",
        description="Create or update an entity (node) in the global Knowledge Graph.",
        inputSchema={
            "type": "object",
            "properties": {
                "id": {"type": "string", "description": "Unique identifier for the entity (e.g. 'Phase56')"},
                "type": {"type": "string", "description": "Type of entity (e.g. 'feature', 'agent', 'concept')"},
                "attributes": {"type": "object", "description": "Key-value pairs of metadata"}
            },
            "required": ["id"]
        }
    ),
    Tool(
        name="kg_create_relation",
        description="Create a directed relationship (edge) between two existing entities in the Knowledge Graph.",
        inputSchema={
            "type": "object",
            "properties": {
                "source": {"type": "string", "description": "ID of the source entity"},
                "target": {"type": "string", "description": "ID of the target entity"},
                "relation_type": {"type": "string", "description": "Type of relationship (e.g. 'depends_on', 'built_by')"}
            },
            "required": ["source", "target"]
        }
    ),
    Tool(
        name="kg_query",
        description="Query the Knowledge Graph. Provide an entity_id to get its details and relations, or omit to get a global summary.",
        inputSchema={
            "type": "object",
            "properties": {
                "entity_id": {"type": "string", "description": "Optional ID of specific entity to query."}
            }
        }
    )
]

HANDLERS = {
    "kg_create_entity": _handle_create_entity,
    "kg_create_relation": _handle_create_relation,
    "kg_query": _handle_query_graph,
}
