"""
Octopoda MCP Server
====================
Model Context Protocol server that exposes Octopoda memory operations as tools.
Any MCP-compatible AI (Claude, Cursor, ChatGPT, VS Code, etc.) can use this.

Setup (Claude Desktop):
    Add to claude_desktop_config.json:
    {
        "mcpServers": {
            "octopoda": {
                "command": "octopoda-mcp"
            }
        }
    }

Run standalone:
    octopoda-mcp
    python -m synrix_runtime.api.mcp_server
"""

import json
import time
from collections import OrderedDict
from mcp.server.fastmcp import FastMCP

mcp = FastMCP("Octopoda Memory")

# -----------------------------------------------------------------------
# Runtime cache (LRU, max 100 agents)
# -----------------------------------------------------------------------

_runtimes: OrderedDict = OrderedDict()
_MAX_RUNTIMES = 100


def _get_runtime(agent_id: str):
    """Get or create an AgentRuntime for the given agent."""
    if agent_id in _runtimes:
        _runtimes.move_to_end(agent_id)
        return _runtimes[agent_id]

    while len(_runtimes) >= _MAX_RUNTIMES:
        _, old_rt = _runtimes.popitem(last=False)
        try:
            old_rt.shutdown()
        except Exception:
            pass

    from synrix_runtime.api.runtime import AgentRuntime
    rt = AgentRuntime(agent_id, agent_type="mcp")
    _runtimes[agent_id] = rt
    return rt


def _parse_value(value: str):
    """Parse a value string as JSON if possible, otherwise return as-is."""
    try:
        return json.loads(value)
    except (json.JSONDecodeError, TypeError):
        return {"value": value}


# -----------------------------------------------------------------------
# Memory Tools
# -----------------------------------------------------------------------

@mcp.tool()
def octopoda_remember(agent_id: str, key: str, value: str, tags: list[str] | None = None) -> dict:
    """Store a persistent memory for an AI agent. The memory survives crashes and restarts.

    Args:
        agent_id: Unique identifier for the agent (e.g. "research_bot", "code_assistant")
        key: Memory key (e.g. "user_preference", "task:current")
        value: Data to store (JSON string or plain text)
        tags: Optional list of tags for categorization
    """
    rt = _get_runtime(agent_id)
    parsed = _parse_value(value)
    result = rt.remember(key, parsed, tags=tags)
    return {
        "success": result.success,
        "key": key,
        "agent_id": agent_id,
        "latency_us": round(result.latency_us, 1),
    }


@mcp.tool()
def octopoda_recall(agent_id: str, key: str) -> dict:
    """Retrieve a stored memory by key.

    Args:
        agent_id: The agent whose memory to read
        key: The memory key to retrieve
    """
    rt = _get_runtime(agent_id)
    result = rt.recall(key)
    return {
        "found": result.found,
        "key": key,
        "value": result.value,
        "latency_us": round(result.latency_us, 1),
    }


@mcp.tool()
def octopoda_search(agent_id: str, prefix: str, limit: int = 20) -> dict:
    """Search an agent's memories by key prefix.

    Args:
        agent_id: The agent whose memories to search
        prefix: Key prefix to search for (e.g. "task:" finds "task:current", "task:history")
        limit: Maximum number of results (default 20)
    """
    rt = _get_runtime(agent_id)
    result = rt.search(prefix, limit=limit)
    return {
        "count": result.count,
        "items": result.items,
        "latency_us": round(result.latency_us, 1),
    }


# -----------------------------------------------------------------------
# Semantic Search & Temporal Tools
# -----------------------------------------------------------------------

@mcp.tool()
def octopoda_recall_similar(agent_id: str, query: str, limit: int = 10) -> dict:
    """Search an agent's memories by meaning (semantic similarity).
    Finds memories related to the query even if the exact words don't match.
    Requires: pip install octopoda[ai]

    Args:
        agent_id: The agent whose memories to search
        query: Natural language query (e.g. "what food does the user like?")
        limit: Maximum number of results (default 10)
    """
    rt = _get_runtime(agent_id)
    result = rt.recall_similar(query, limit=limit)
    return {
        "count": result.count,
        "items": result.items,
        "latency_us": round(result.latency_us, 1),
    }


@mcp.tool()
def octopoda_recall_history(agent_id: str, key: str) -> dict:
    """Get the full timeline of how a memory changed over time.
    Shows all versions with timestamps for when each was valid.

    Args:
        agent_id: The agent whose memory to inspect
        key: The memory key to get history for
    """
    rt = _get_runtime(agent_id)
    result = rt.recall_history(key)
    return {
        "key": result.key,
        "current_version": result.current_version,
        "versions": result.versions,
        "latency_us": round(result.latency_us, 1),
    }


@mcp.tool()
def octopoda_related(agent_id: str, entity: str) -> dict:
    """Query the knowledge graph for an entity and its connections.
    Shows what other entities are related and how.
    Requires: pip install octopoda[nlp]

    Args:
        agent_id: The agent whose knowledge graph to query
        entity: Entity name to look up (e.g. "London", "Alice")
    """
    rt = _get_runtime(agent_id)
    result = rt.related(entity)
    return {
        "found": result.found,
        "entity": result.entity,
        "entity_type": result.entity_type,
        "relationships": result.relationships,
        "latency_us": round(result.latency_us, 1),
    }


# -----------------------------------------------------------------------
# Snapshot Tools
# -----------------------------------------------------------------------

@mcp.tool()
def octopoda_snapshot(agent_id: str, label: str | None = None) -> dict:
    """Take a snapshot (checkpoint) of all agent memory. Use before risky operations.

    Args:
        agent_id: The agent whose memory to snapshot
        label: Optional label for the snapshot (auto-generated if omitted)
    """
    rt = _get_runtime(agent_id)
    result = rt.snapshot(label)
    return {
        "label": result.label,
        "keys_captured": result.keys_captured,
        "size_bytes": result.size_bytes,
        "latency_us": round(result.latency_us, 1),
    }


@mcp.tool()
def octopoda_restore(agent_id: str, label: str | None = None) -> dict:
    """Restore agent memory from a snapshot. Reverts to the saved state.

    Args:
        agent_id: The agent whose memory to restore
        label: Snapshot label to restore from (latest if omitted)
    """
    rt = _get_runtime(agent_id)
    result = rt.restore(label)
    return {
        "label": result.label,
        "keys_restored": result.keys_restored,
        "recovery_time_us": round(result.recovery_time_us, 1),
    }


# -----------------------------------------------------------------------
# Shared Memory Tools
# -----------------------------------------------------------------------

@mcp.tool()
def octopoda_share(agent_id: str, key: str, value: str, space: str = "global") -> dict:
    """Write to shared memory that other agents can read.

    Args:
        agent_id: The agent writing the data
        key: Shared memory key
        value: Data to share (JSON string or plain text)
        space: Memory space name (default "global")
    """
    rt = _get_runtime(agent_id)
    parsed = _parse_value(value)
    result = rt.share(key, parsed, space=space)
    return {
        "success": result.success,
        "key": key,
        "space": space,
        "latency_us": round(result.latency_us, 1),
    }


@mcp.tool()
def octopoda_read_shared(agent_id: str, key: str, space: str = "global") -> dict:
    """Read from shared memory written by any agent.

    Args:
        agent_id: The agent reading the data
        key: Shared memory key to read
        space: Memory space name (default "global")
    """
    rt = _get_runtime(agent_id)
    result = rt.read_shared(key, space=space)
    return {
        "found": result.found,
        "key": key,
        "space": space,
        "value": result.value,
        "latency_us": round(result.latency_us, 1),
    }


# -----------------------------------------------------------------------
# Agent Management Tools
# -----------------------------------------------------------------------

@mcp.tool()
def octopoda_list_agents() -> dict:
    """List all registered agents and their current state."""
    try:
        from synrix_runtime.core.daemon import RuntimeDaemon
        daemon = RuntimeDaemon.get_instance()
        if daemon.running:
            agents = daemon.get_active_agents()
            return {
                "count": len(agents),
                "agents": [
                    {
                        "agent_id": a.get("agent_id"),
                        "type": a.get("type", "unknown"),
                        "state": a.get("state", "unknown"),
                    }
                    for a in agents
                ],
            }
    except Exception:
        pass

    # Fallback: list cached runtimes
    return {
        "count": len(_runtimes),
        "agents": [
            {"agent_id": aid, "type": "mcp", "state": "running"}
            for aid in _runtimes
        ],
    }


@mcp.tool()
def octopoda_agent_stats(agent_id: str) -> dict:
    """Get performance statistics for an agent.

    Args:
        agent_id: The agent to get stats for
    """
    rt = _get_runtime(agent_id)
    stats = rt.get_stats()
    return {
        "agent_id": stats.agent_id,
        "total_operations": stats.total_operations,
        "total_writes": stats.total_writes,
        "total_reads": stats.total_reads,
        "total_queries": stats.total_queries,
        "avg_write_latency_us": round(stats.avg_write_latency_us, 1),
        "avg_read_latency_us": round(stats.avg_read_latency_us, 1),
        "performance_score": round(stats.performance_score, 1),
        "uptime_seconds": round(stats.uptime_seconds, 1),
    }


# -----------------------------------------------------------------------
# Audit Tool
# -----------------------------------------------------------------------

@mcp.tool()
def octopoda_log_decision(
    agent_id: str, decision: str, reasoning: str, context: str | None = None
) -> dict:
    """Log an agent decision with full audit trail and memory snapshot.

    Args:
        agent_id: The agent making the decision
        decision: What was decided
        reasoning: Why this decision was made
        context: Optional JSON string with additional context
    """
    rt = _get_runtime(agent_id)
    ctx = _parse_value(context) if context else {}
    rt.log_decision(decision, reasoning, ctx)
    return {
        "agent_id": agent_id,
        "logged": True,
        "decision": decision,
    }


# -----------------------------------------------------------------------
# Entry point
# -----------------------------------------------------------------------

def main():
    """Run the MCP server (stdio transport)."""
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
