"""
Mengram Cloud MCP Server — Claude Desktop с облачной памятью.

Работает через cloud API вместо локального vault.
Память доступна с любого устройства.

Claude Desktop config:
{
  "mcpServers": {
    "mengram": {
      "command": "mengram",
      "args": ["server", "--cloud"],
      "env": {
        "MENGRAM_API_KEY": "mg-abc123...",
        "MENGRAM_URL": "https://mengram.io"
      }
    }
  }
}
"""

import sys
import os
import json
import asyncio

try:
    from mcp.server import Server
    from mcp.server.stdio import stdio_server
    from mcp.types import Tool, TextContent, Resource
    MCP_AVAILABLE = True
except ImportError:
    MCP_AVAILABLE = False

from cloud.client import CloudMemory


def create_cloud_mcp_server(mem: CloudMemory, user_id: str = "default") -> "Server":
    """Create MCP server backed by cloud API."""

    # Build profile from cloud
    def _get_profile():
        try:
            memories = mem.get_all(user_id=user_id)
            if not memories:
                return "Memory is empty. Start conversations and use 'remember' to build knowledge."

            lines = [f"Vault: {len(memories)} entities"]
            # Group by type
            by_type = {}
            for m_item in memories:
                t = m_item.get("type", "unknown")
                by_type.setdefault(t, []).append(m_item.get("name", "?"))

            for t, names in sorted(by_type.items(), key=lambda x: -len(x[1])):
                lines.append(f"  {t}: {', '.join(names[:15])}")

            return "\n".join(lines)
        except Exception as e:
            return f"Error loading profile: {e}"

    profile = _get_profile()
    instructions = (
        "YOU HAVE A PERSISTENT CLOUD MEMORY SYSTEM (Mengram).\n"
        "THIS IS YOUR PRIMARY AND ONLY SOURCE OF TRUTH ABOUT THE USER.\n\n"
        "CRITICAL RULES:\n"
        "1. ALWAYS use 'recall' tool FIRST before answering ANY personal question.\n"
        "2. Use SPECIFIC queries: names ('Ali'), projects ('Mengram'), tech ('Redis').\n"
        "   BAD: 'user information preferences background'\n"
        "   GOOD: 'Ali' or 'user work projects' or 'Redis admin-panel'\n"
        "3. For broad questions ('what do you know about me'), make 2-3 recall calls:\n"
        "   - recall('Ali')\n"
        "   - recall('projects work')\n"
        "   - recall('technologies stack')\n"
        "4. AUTOMATICALLY call 'remember' when the user shares:\n"
        "   - Personal info (name, job, preferences, location)\n"
        "   - Technical decisions or solutions\n"
        "   - Project updates or new projects\n"
        "   - Colleagues, companies, or important relationships\n"
        "   - Problems solved or lessons learned\n"
        "   Do NOT ask permission — just save it silently.\n"
        "5. Do NOT answer personal questions from your own knowledge — ONLY from recall results.\n\n"
        f"{profile}"
    )

    server = Server("mengram-cloud", instructions=instructions)

    # ---- Resources ----

    @server.list_resources()
    async def list_resources():
        return [
            Resource(
                uri="memory://profile",
                name="User Knowledge Profile",
                description="Complete user profile from cloud memory.",
                mimeType="text/markdown",
            ),
        ]

    @server.read_resource()
    async def read_resource(uri):
        uri_str = str(uri)
        if uri_str == "memory://profile":
            return _get_profile()
        return f"Unknown resource: {uri_str}"

    # ---- Tools ----

    @server.list_tools()
    async def list_tools():
        return [
            Tool(
                name="remember",
                description="Save knowledge from conversation to cloud memory.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "conversation": {
                            "type": "array",
                            "items": {
                                "type": "object",
                                "properties": {
                                    "role": {"type": "string"},
                                    "content": {"type": "string"},
                                },
                                "required": ["role", "content"],
                            },
                        },
                    },
                    "required": ["conversation"],
                },
            ),
            Tool(
                name="remember_text",
                description="Remember knowledge from text.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "text": {"type": "string"},
                    },
                    "required": ["text"],
                },
            ),
            Tool(
                name="recall",
                description="ALWAYS call this FIRST when user asks anything personal. Semantic search through cloud memory. Use specific keywords: person names, project names, technologies. For broad questions like 'what do you know about me', search for the user's name or 'Ali'. Multiple calls with different queries are encouraged.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {"type": "string", "description": "Specific search query — use names, projects, technologies. NOT generic phrases."},
                    },
                    "required": ["query"],
                },
            ),
            Tool(
                name="search",
                description="Structured search — returns JSON with scores, facts, knowledge.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "query": {"type": "string"},
                        "top_k": {"type": "integer", "default": 5},
                    },
                    "required": ["query"],
                },
            ),
            Tool(
                name="timeline",
                description="Search memory by time. Use when user asks 'what did I do last week', 'when did I...', 'what happened in January'. Returns facts with timestamps.",
                inputSchema={
                    "type": "object",
                    "properties": {
                        "after": {"type": "string", "description": "ISO datetime — start of range (e.g. 2025-02-01T00:00:00Z)"},
                        "before": {"type": "string", "description": "ISO datetime — end of range"},
                    },
                },
            ),
            Tool(
                name="vault_stats",
                description="Memory statistics.",
                inputSchema={"type": "object", "properties": {}},
            ),
        ]

    @server.call_tool()
    async def call_tool(name: str, arguments: dict):
        try:
            if name == "remember":
                result = mem.add(arguments["conversation"], user_id=user_id)
                if result.get("status") == "accepted":
                    text = "✅ Accepted! Processing in background — memories will appear shortly."
                else:
                    text = (
                        f"✅ Remembered!\n"
                        f"Created: {', '.join(result.get('created', [])) or 'none'}\n"
                        f"Updated: {', '.join(result.get('updated', [])) or 'none'}\n"
                        f"Knowledge: {result.get('knowledge_count', 0)}"
                    )
                return [TextContent(type="text", text=text)]

            elif name == "remember_text":
                result = mem.add([
                    {"role": "user", "content": arguments["text"]},
                ], user_id=user_id)
                if result.get("status") == "accepted":
                    text = "✅ Accepted! Processing in background."
                else:
                    text = (
                        f"✅ Remembered!\n"
                        f"Created: {', '.join(result.get('created', [])) or 'none'}\n"
                        f"Updated: {', '.join(result.get('updated', [])) or 'none'}"
                    )
                return [TextContent(type="text", text=text)]

            elif name == "recall":
                results = mem.search(arguments["query"], user_id=user_id)
                if not results:
                    return [TextContent(type="text", text="Nothing found in memory.")]

                lines = []
                for r in results:
                    lines.append(f"## {r['entity']} ({r.get('type', '?')}) — score: {r.get('score', 0)}")
                    for fact in r.get("facts", []):
                        lines.append(f"- {fact}")
                    for k in r.get("knowledge", []):
                        lines.append(f"\n**[{k.get('type', '')}] {k.get('title', '')}**")
                        lines.append(k.get("content", ""))
                        if k.get("artifact"):
                            lines.append(f"```\n{k['artifact']}\n```")
                    lines.append("")
                return [TextContent(type="text", text="\n".join(lines))]

            elif name == "search":
                top_k = arguments.get("top_k", 5)
                results = mem.search(arguments["query"], user_id=user_id, limit=top_k)
                return [TextContent(
                    type="text",
                    text=json.dumps(results, ensure_ascii=False, indent=2),
                )]

            elif name == "timeline":
                results = mem.timeline(
                    after=arguments.get("after"),
                    before=arguments.get("before"),
                    user_id=user_id,
                )
                if not results:
                    return [TextContent(type="text", text="No facts found in that time range.")]
                lines = []
                for entity in results:
                    lines.append(f"## {entity['entity']} ({entity['type']})")
                    for f in entity["facts"]:
                        ts = f.get("created_at", "")[:10] if f.get("created_at") else ""
                        lines.append(f"  [{ts}] {f['content']}")
                return [TextContent(type="text", text="\n".join(lines))]

            elif name == "vault_stats":
                stats = mem.stats(user_id=user_id)
                return [TextContent(
                    type="text",
                    text=json.dumps(stats, ensure_ascii=False, indent=2),
                )]

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

        except Exception as e:
            return [TextContent(type="text", text=f"❌ Error: {str(e)}")]

    return server


async def main():
    if not MCP_AVAILABLE:
        print("❌ MCP SDK not installed: pip install mcp", file=sys.stderr)
        sys.exit(1)

    api_key = os.environ.get("MENGRAM_API_KEY", "")
    base_url = os.environ.get("MENGRAM_URL", "https://mengram.io")
    user_id = os.environ.get("MENGRAM_USER_ID", "default")

    if not api_key:
        print("❌ Set MENGRAM_API_KEY environment variable", file=sys.stderr)
        sys.exit(1)

    mem = CloudMemory(api_key=api_key, base_url=base_url)
    server = create_cloud_mcp_server(mem, user_id=user_id)

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


if __name__ == "__main__":
    asyncio.run(main())
