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

All memory operations go through the Octopoda Cloud API (api.octopodas.com).
Requires OCTOPODA_API_KEY — get your free key at https://octopodas.com

Setup (Claude Desktop):
    Add to claude_desktop_config.json:
    {
        "mcpServers": {
            "octopoda": {
                "command": "octopoda-mcp",
                "env": {
                    "OCTOPODA_API_KEY": "sk-octopoda-YOUR_KEY"
                }
            }
        }
    }

Install:
    pip install octopoda[mcp]

Run standalone:
    OCTOPODA_API_KEY=sk-octopoda-... octopoda-mcp
    OCTOPODA_API_KEY=sk-octopoda-... python -m synrix.mcp_server
"""

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

mcp = FastMCP("Octopoda Memory")

# -----------------------------------------------------------------------
# Cloud client cache (LRU, max 100 agents)
# -----------------------------------------------------------------------

_client = None
_agents: OrderedDict = OrderedDict()
_MAX_AGENTS = 100


def _get_client():
    """Get or create the Octopoda cloud client (singleton)."""
    global _client
    if _client is not None:
        return _client

    from synrix.cloud import Octopoda
    _client = Octopoda()  # reads OCTOPODA_API_KEY from env
    return _client


def _get_agent(agent_id: str):
    """Get or create a cloud Agent handle for the given agent_id."""
    if agent_id in _agents:
        _agents.move_to_end(agent_id)
        return _agents[agent_id]

    while len(_agents) >= _MAX_AGENTS:
        _agents.popitem(last=False)

    client = _get_client()
    agent = client.agent(agent_id, metadata={"type": "mcp"})
    _agents[agent_id] = agent
    return agent


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}


def _timed(fn):
    """Call fn, return (result, latency_us)."""
    t0 = time.perf_counter()
    result = fn()
    latency = (time.perf_counter() - t0) * 1_000_000
    return result, round(latency, 1)


# -----------------------------------------------------------------------
# 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. Memory is stored in the cloud and persists across sessions.

    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
    """
    agent = _get_agent(agent_id)
    parsed = _parse_value(value)
    result, latency = _timed(lambda: agent.write(key, parsed, tags=tags))
    return {
        "success": True,
        "key": key,
        "agent_id": agent_id,
        "latency_us": latency,
    }


@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
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.read(key))
    return {
        "found": result is not None,
        "key": key,
        "value": result,
        "latency_us": latency,
    }


@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)
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.keys(prefix=prefix, limit=limit))
    return {
        "count": len(result),
        "items": result,
        "latency_us": latency,
    }


# -----------------------------------------------------------------------
# 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.

    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)
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.search(query, limit=limit))
    return {
        "count": len(result),
        "items": result,
        "latency_us": latency,
    }


@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
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.history(key))
    return {
        "key": key,
        "versions": result,
        "latency_us": latency,
    }


@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.

    Args:
        agent_id: The agent whose knowledge graph to query
        entity: Entity name to look up (e.g. "London", "Alice")
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.related(entity))
    return {
        "entity": entity,
        "relationships": result,
        "latency_us": latency,
    }


# -----------------------------------------------------------------------
# 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)
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.snapshot(label))
    return {
        "label": result.get("label", label),
        "keys_captured": result.get("keys_captured", 0),
        "latency_us": latency,
    }


@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)
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.restore(label or "latest"))
    return {
        "label": label,
        "keys_restored": result.get("keys_restored", 0),
        "latency_us": latency,
    }


# -----------------------------------------------------------------------
# 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")
    """
    agent = _get_agent(agent_id)
    parsed = _parse_value(value)
    result, latency = _timed(lambda: agent.share(space, key, parsed))
    return {
        "success": True,
        "key": key,
        "space": space,
        "latency_us": latency,
    }


@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")
    """
    client = _get_client()
    result, latency = _timed(lambda: client.read_shared(space, key))
    return {
        "found": result is not None and "error" not in (result or {}),
        "key": key,
        "space": space,
        "value": result,
        "latency_us": latency,
    }


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

@mcp.tool()
def octopoda_list_agents() -> dict:
    """List all registered agents in your Octopoda account."""
    client = _get_client()
    result, latency = _timed(lambda: client.agents())
    return {
        "count": len(result),
        "agents": result,
        "latency_us": latency,
    }


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

    Args:
        agent_id: The agent to get stats for
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.metrics())
    return {
        "agent_id": agent_id,
        "metrics": result,
        "latency_us": latency,
    }


# -----------------------------------------------------------------------
# Conversation Processing Tools
# -----------------------------------------------------------------------

@mcp.tool()
def octopoda_process_conversation(agent_id: str, user_message: str, assistant_response: str) -> dict:
    """Process a conversation turn — automatically extracts and stores memories.
    Call this after your agent responds to learn from the conversation.

    Args:
        agent_id: The agent to store memories for
        user_message: What the user said
        assistant_response: What the assistant replied
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.process_conversation(
        messages=[
            {"role": "user", "content": user_message},
            {"role": "assistant", "content": assistant_response},
        ],
        extract_preferences=True,
        extract_facts=True,
        extract_decisions=True,
        namespace="conversations",
    ))
    return {
        "agent_id": agent_id,
        "processed": True,
        "details": result,
        "latency_us": latency,
    }


@mcp.tool()
def octopoda_get_context(agent_id: str, query: str, limit: int = 10) -> dict:
    """Get relevant context from memory before generating a response.
    Returns memories related to the query, ready to inject into your prompt.

    Args:
        agent_id: The agent whose memories to search
        query: The current user message or topic
        limit: Max memories to retrieve (default 10)
    """
    agent = _get_agent(agent_id)
    result, latency = _timed(lambda: agent.get_context(query, limit=limit, format="text"))
    return {
        "context": result.get("context", ""),
        "latency_us": latency,
    }


# -----------------------------------------------------------------------
# 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.

    Args:
        agent_id: The agent making the decision
        decision: What was decided ("allow", "deny", or "escalate")
        reasoning: Why this decision was made
        context: Optional JSON string with additional context
    """
    agent = _get_agent(agent_id)
    ctx = _parse_value(context) if context else {}
    result, latency = _timed(lambda: agent.decide(decision, reasoning, ctx))
    return {
        "agent_id": agent_id,
        "logged": True,
        "decision": decision,
        "latency_us": latency,
    }


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

def main():
    """Run the MCP server (stdio transport)."""
    api_key = os.environ.get("OCTOPODA_API_KEY", "")
    if not api_key:
        import sys
        print("ERROR: OCTOPODA_API_KEY not set.", file=sys.stderr)
        print("Get your free key at https://octopodas.com", file=sys.stderr)
        print("Then: OCTOPODA_API_KEY=sk-octopoda-... octopoda-mcp", file=sys.stderr)
        sys.exit(1)
    mcp.run(transport="stdio")


if __name__ == "__main__":
    main()
