#!/usr/bin/env python3
"""
Mnemo MCP Server
Provides memory capabilities to Claude and other MCP-compatible LLMs

Usage:
    uvx mnemo-memory  # Run as MCP server
    
Or in Claude config:
    {
        "mcpServers": {
            "mnemo": {"command": "uvx", "args": ["mnemo-memory"]}
        }
    }
"""

import json
import sys
from typing import Any, Dict, List

try:
    from mcp.server import Server
    from mcp.types import Tool, TextContent
    HAS_MCP = True
except ImportError:
    HAS_MCP = False
    print("Warning: MCP not installed. Install with: pip install mcp", file=sys.stderr)

from mnemo import Mnemo, should_inject_memory

# Global memory instance
memory = Mnemo(use_real_embeddings=True)

# Tool definitions
TOOLS = [
    {
        "name": "mnemo_add",
        "description": "Store a new memory. Use this to save important information for later retrieval.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "content": {
                    "type": "string",
                    "description": "The content to remember"
                },
                "metadata": {
                    "type": "object",
                    "description": "Optional metadata (tags, source, etc.)"
                }
            },
            "required": ["content"]
        }
    },
    {
        "name": "mnemo_search",
        "description": "Search stored memories. Returns relevant memories ranked by similarity.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query"
                },
                "top_k": {
                    "type": "integer",
                    "description": "Number of results (default: 5)",
                    "default": 5
                }
            },
            "required": ["query"]
        }
    },
    {
        "name": "mnemo_should_inject",
        "description": "Check if memory should be injected for this query. Uses context-check algorithm with 90% accuracy.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The query to check"
                },
                "context": {
                    "type": "string",
                    "description": "Optional additional context"
                }
            },
            "required": ["query"]
        }
    },
    {
        "name": "mnemo_get_context",
        "description": "Get formatted memory context for injection into prompts.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query for finding relevant context"
                },
                "top_k": {
                    "type": "integer",
                    "description": "Number of memories to include (default: 3)",
                    "default": 3
                }
            },
            "required": ["query"]
        }
    },
    {
        "name": "mnemo_list",
        "description": "List all stored memories.",
        "inputSchema": {
            "type": "object",
            "properties": {}
        }
    },
    {
        "name": "mnemo_delete",
        "description": "Delete a specific memory by ID.",
        "inputSchema": {
            "type": "object",
            "properties": {
                "memory_id": {
                    "type": "string",
                    "description": "ID of the memory to delete"
                }
            },
            "required": ["memory_id"]
        }
    },
    {
        "name": "mnemo_stats",
        "description": "Get memory system statistics.",
        "inputSchema": {
            "type": "object",
            "properties": {}
        }
    },
    {
        "name": "mnemo_clear",
        "description": "Clear all stored memories. Use with caution!",
        "inputSchema": {
            "type": "object",
            "properties": {
                "confirm": {
                    "type": "boolean",
                    "description": "Must be true to confirm deletion"
                }
            },
            "required": ["confirm"]
        }
    }
]


def handle_tool_call(name: str, arguments: Dict[str, Any]) -> str:
    """Handle a tool call and return result"""
    
    if name == "mnemo_add":
        content = arguments.get("content", "")
        metadata = arguments.get("metadata", {})
        mem_id = memory.add(content, metadata)
        return json.dumps({
            "status": "success",
            "memory_id": mem_id,
            "message": f"Memory stored successfully"
        })
    
    elif name == "mnemo_search":
        query = arguments.get("query", "")
        top_k = arguments.get("top_k", 5)
        results = memory.search(query, top_k=top_k)
        
        return json.dumps({
            "status": "success",
            "count": len(results),
            "results": [
                {
                    "id": r.id,
                    "content": r.content,
                    "score": round(r.score, 3),
                    "metadata": r.metadata
                }
                for r in results
            ]
        })
    
    elif name == "mnemo_should_inject":
        query = arguments.get("query", "")
        context = arguments.get("context", "")
        should, reason = should_inject_memory(query, context)
        
        return json.dumps({
            "should_inject": should,
            "reason": reason,
            "recommendation": "Inject memory context" if should else "Skip memory - standalone query"
        })
    
    elif name == "mnemo_get_context":
        query = arguments.get("query", "")
        top_k = arguments.get("top_k", 3)
        context = memory.get_context(query, top_k=top_k)
        
        return json.dumps({
            "status": "success",
            "context": context if context else None,
            "message": "Context retrieved" if context else "No relevant context found"
        })
    
    elif name == "mnemo_list":
        memories = memory.list_all()
        return json.dumps({
            "status": "success",
            "count": len(memories),
            "memories": [
                {
                    "id": m.id,
                    "content": m.content[:200] + "..." if len(m.content) > 200 else m.content,
                    "created_at": m.created_at,
                    "metadata": m.metadata
                }
                for m in memories
            ]
        })
    
    elif name == "mnemo_delete":
        memory_id = arguments.get("memory_id", "")
        success = memory.delete(memory_id)
        
        return json.dumps({
            "status": "success" if success else "error",
            "message": f"Memory {memory_id} deleted" if success else f"Memory {memory_id} not found"
        })
    
    elif name == "mnemo_stats":
        stats = memory.get_stats()
        return json.dumps({
            "status": "success",
            "stats": stats
        })
    
    elif name == "mnemo_clear":
        if arguments.get("confirm", False):
            memory.clear()
            return json.dumps({
                "status": "success",
                "message": "All memories cleared"
            })
        else:
            return json.dumps({
                "status": "error",
                "message": "Must set confirm=true to clear all memories"
            })
    
    else:
        return json.dumps({
            "status": "error",
            "message": f"Unknown tool: {name}"
        })


def run_mcp_server():
    """Run as MCP server"""
    if not HAS_MCP:
        print("Error: MCP not installed. Install with: pip install mcp", file=sys.stderr)
        sys.exit(1)
    
    server = Server("mnemo-memory")
    
    @server.list_tools()
    async def list_tools() -> List[Tool]:
        return [
            Tool(
                name=tool["name"],
                description=tool["description"],
                inputSchema=tool["inputSchema"]
            )
            for tool in TOOLS
        ]
    
    @server.call_tool()
    async def call_tool(name: str, arguments: Dict[str, Any]) -> List[TextContent]:
        result = handle_tool_call(name, arguments)
        return [TextContent(type="text", text=result)]
    
    # Run server
    import asyncio
    asyncio.run(server.run())


def run_cli():
    """Simple CLI for testing"""
    print("Mnemo Memory System - CLI Mode")
    print("Commands: add, search, inject, context, list, stats, clear, quit")
    print("-" * 50)
    
    while True:
        try:
            cmd = input("\n> ").strip().lower()
            
            if cmd == "quit":
                break
            elif cmd == "add":
                content = input("Content: ")
                result = handle_tool_call("mnemo_add", {"content": content})
                print(result)
            elif cmd == "search":
                query = input("Query: ")
                result = handle_tool_call("mnemo_search", {"query": query})
                print(result)
            elif cmd == "inject":
                query = input("Query: ")
                result = handle_tool_call("mnemo_should_inject", {"query": query})
                print(result)
            elif cmd == "context":
                query = input("Query: ")
                result = handle_tool_call("mnemo_get_context", {"query": query})
                print(result)
            elif cmd == "list":
                result = handle_tool_call("mnemo_list", {})
                print(result)
            elif cmd == "stats":
                result = handle_tool_call("mnemo_stats", {})
                print(result)
            elif cmd == "clear":
                confirm = input("Are you sure? (yes/no): ")
                if confirm.lower() == "yes":
                    result = handle_tool_call("mnemo_clear", {"confirm": True})
                    print(result)
            else:
                print("Unknown command. Use: add, search, inject, context, list, stats, clear, quit")
        
        except KeyboardInterrupt:
            print("\nBye!")
            break
        except Exception as e:
            print(f"Error: {e}")


if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "--cli":
        run_cli()
    else:
        run_mcp_server()