"""Agent Memory MCP Server — V12

Exposes AgentMemory + Spirit as MCP tools for any MCP-compatible Agent.
Supports both official MCP SDK (stdio/streamable-http) and HTTP fallback.
"""
from __future__ import annotations

import json
import logging
import time
import hashlib
from typing import Any

logger = logging.getLogger(__name__)

# Try importing the official MCP SDK
_MCP_AVAILABLE = False
try:
    from mcp.server.fastmcp import FastMCP
    _MCP_AVAILABLE = True
except ImportError:
    logger.info("MCP SDK not installed. HTTP fallback will be used. Install with: pip install 'mcp[cli]'")

from .mcp_config import MCPConfig


class _RateTracker:
    """Simple rate tracker for MCP requests."""
    def __init__(self, max_per_minute: int = 60):
        self.max_per_minute = max_per_minute
        self._timestamps: list[float] = []

    def check(self) -> bool:
        """Returns True if request is allowed, False if rate limited."""
        now = time.time()
        self._timestamps = [t for t in self._timestamps if now - t < 60]
        if len(self._timestamps) >= self.max_per_minute:
            return False
        self._timestamps.append(now)
        return True


class AgentMemoryMCPServer:
    """MCP Server for Agent Memory — any Agent can plug in and use.

    Exposes 11 MCP tools:
    - memory.remember: Store a memory (auto-tags source agent)
    - memory.recall: Retrieve memories (filtered by agent identity)
    - memory.spirit_check: Spirit proactive check
    - memory.get_profile: Get user profile (privacy-filtered)
    - memory.report: Generate memory report
    - memory.share_skill: Share a skill with other agents
    - memory.learn_skill: Learn a skill from other agents
    - memory.command: Natural language Spirit command
    """

    def __init__(self, memory, spirit=None, agent_manager=None, config: MCPConfig | None = None):
        self.memory = memory
        self.spirit = spirit
        self.agent_manager = agent_manager
        self.config = config or MCPConfig()

        self._read_tracker = _RateTracker(self.config.max_requests_per_minute)
        self._write_tracker = _RateTracker(self.config.max_write_per_minute)

        self._mcp_server = None
        if _MCP_AVAILABLE:
            self._mcp_server = FastMCP(
                self.config.server_name,
                version=self.config.server_version,
            )
            self._register_tools()

    def _check_auth(self, api_key: str | None = None) -> bool:
        """Check authentication. Supports both API key and Agent Token."""
        if not self.config.require_auth:
            return True
        if not api_key:
            return False
        # First try AgentManager token auth
        if self.agent_manager:
            agent = self.agent_manager.authenticate_agent(api_key)
            if agent:
                return True
        # Fallback to static API key
        return api_key == self.config.api_key

    def _resolve_agent_id(self, source_agent: str, api_key: str | None = None) -> str:
        """Resolve actual agent_id from auth token, falling back to declared source_agent."""
        if self.agent_manager and api_key:
            agent = self.agent_manager.authenticate_agent(api_key)
            if agent:
                return agent["agent_id"]
        return source_agent

    def _check_rate(self, is_write: bool = False) -> bool:
        if is_write:
            return self._write_tracker.check()
        return self._read_tracker.check()

    def _handle_error(self, error: Exception, context: str = "") -> dict[str, Any]:
        """Convert exceptions to user-friendly error responses."""
        import traceback
        error_type = type(error).__name__
        error_msg = str(error)

        if "not available" in error_msg.lower() or "not_initialized" in error_msg.lower():
            return {
                "status": "error",
                "error": error_msg,
                "error_type": error_type,
            }

        generic_messages = {
            "MemoryPermissionError": "权限不足，无法执行此操作",
            "PermissionDeniedError": "权限不足，无法执行此操作",
            "ValueError": "参数值无效，请检查输入",
            "TimeoutError": "操作超时，请重试",
            "ConnectionError": "连接失败，请检查网络",
            "MemoryError": "内存不足，请清理记忆或减少查询数量",
        }

        if error_type in generic_messages:
            friendly_msg = generic_messages[error_type]
        else:
            friendly_msg = f"操作失败：{error_type}"

        logger.debug("Error in %s: %s\n%s", context, error, traceback.format_exc())

        return {
            "status": "error",
            "error": friendly_msg,
            "error_type": error_type,
        }

    # --- Tool implementations ---

    def tool_remember(self, content: str, source_agent: str = "unknown",
                      sensitivity: str | None = None, scope: str | None = None,
                      importance: str = "medium", significance: str = "medium",
                      topics: list[str] | None = None,
                      api_key: str | None = None) -> dict[str, Any]:
        """Store a memory. Spirit auto-tags source agent and manages dedup."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed", "message": "API Key 无效"}
        if not self._check_rate(is_write=True):
            return {"error": "rate_limited", "message": "写入频率超限"}

        sens = sensitivity or self.config.default_sensitivity
        mem_scope = scope or self.config.default_scope

        try:
            memory_id = self.memory.remember(
                content=content,
                importance=importance,
                significance=significance,
                source=f"mcp:{source_agent}",
            )
            return {
                "status": "stored",
                "memory_id": memory_id,
                "source_agent": source_agent,
                "sensitivity": sens,
                "scope": mem_scope,
            }
        except Exception as e:
            return self._handle_error(e, "tool_remember")

    def tool_recall(self, query: str, source_agent: str = "unknown",
                    top_k: int = 10, api_key: str | None = None) -> dict[str, Any]:
        """Retrieve memories. Privacy-filtered based on agent identity."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed", "message": "API Key 无效"}
        if not self._check_rate(is_write=False):
            return {"error": "rate_limited", "message": "读取频率超限"}

        # Resolve agent context for privacy filtering
        resolved_agent = self._resolve_agent_id(source_agent, api_key)
        agent_ctx = None
        if self.agent_manager:
            agent_ctx = self.agent_manager.get_agent_context(resolved_agent)

        # Build privacy filters for index-level pushdown
        recall_kwargs = {"query": query, "top_k": top_k}
        if agent_ctx:
            recall_kwargs["tenant_id"] = agent_ctx.get("agent_scope", "")
            # Only filter by visibility if the agent has limited access
            max_vis = agent_ctx.get("max_sensitivity", "team")
            if max_vis != "private":
                recall_kwargs["visibility_filter"] = ["public", "team"] if max_vis == "team" else ["public"]

        try:
            results = self.memory.recall(**recall_kwargs)
            # Basic filtering: convert to safe dicts
            memories = []
            if isinstance(results, dict):
                for mid, mem in results.items():
                    if isinstance(mem, dict):
                        memories.append({
                            "memory_id": mid,
                            "content": mem.get("content", ""),
                            "importance": mem.get("importance", "medium"),
                            "topics": mem.get("topics", []),
                            "timestamp": mem.get("time_ts", 0),
                        })
            elif isinstance(results, list):
                for mem in results:
                    if isinstance(mem, dict):
                        memories.append({
                            "memory_id": mem.get("memory_id", ""),
                            "content": mem.get("content", ""),
                            "importance": mem.get("importance", "medium"),
                            "topics": mem.get("topics", []),
                            "timestamp": mem.get("time_ts", 0),
                        })
            return {
                "status": "ok",
                "query": query,
                "count": len(memories),
                "memories": memories,
            }
        except Exception as e:
            return self._handle_error(e, "tool_recall")

    def tool_spirit_check(self, api_key: str | None = None) -> dict[str, Any]:
        """Spirit proactive check — butler scans for issues and suggests actions."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=False):
            return {"error": "rate_limited"}

        if not self.spirit:
            return {"error": "spirit_not_available", "message": "Spirit 管家未初始化"}

        try:
            result = self.spirit.proactive_check()
            return {"status": "ok", **result}
        except Exception as e:
            logger.error("MCP tool_spirit_check: %s", e)
            return {"error": "spirit_check_failed", "message": str(e)}

    def tool_get_profile(self, source_agent: str = "unknown",
                         api_key: str | None = None) -> dict[str, Any]:
        """Get user profile. Spirit decides how much to reveal based on agent identity."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=False):
            return {"error": "rate_limited"}

        if not self.spirit:
            return {"error": "spirit_not_available"}

        try:
            profile = self.spirit.get_profile()
            return {"status": "ok", "profile": profile}
        except Exception as e:
            logger.error("MCP tool_get_profile: %s", e)
            return {"error": "profile_failed", "message": str(e)}

    def tool_report(self, report_type: str = "daily",
                    api_key: str | None = None) -> dict[str, Any]:
        """Generate memory report (daily/weekly/monthly)."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=False):
            return {"error": "rate_limited"}

        if not self.spirit:
            return {"error": "spirit_not_available"}

        try:
            report = self.spirit.report(report_type=report_type)
            return {"status": "ok", "report_type": report_type, "report": report}
        except Exception as e:
            logger.error("MCP tool_report: %s", e)
            return {"error": "report_failed", "message": str(e)}

    def tool_share_skill(self, skill_name: str, skill_description: str,
                         skill_steps: list[str], category: str = "general",
                         source_agent: str = "unknown",
                         visibility: str = "internal",
                         api_key: str | None = None) -> dict[str, Any]:
        """Share a skill with other agents through the butler."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=True):
            return {"error": "rate_limited"}

        try:
            skill_id = f"skill_{hashlib.sha256(f'{source_agent}:{skill_name}'.encode()).hexdigest()[:12]}"
            # Store skill as a special memory
            skill_content = json.dumps({
                "type": "skill_package",
                "skill_id": skill_id,
                "name": skill_name,
                "description": skill_description,
                "category": category,
                "steps": skill_steps,
                "source_agent": source_agent,
                "visibility": visibility,
                "created_at": time.time(),
            }, ensure_ascii=False)

            memory_id = self.memory.remember(
                content=skill_content,
                importance="medium",
                source=f"mcp_skill:{source_agent}",
            )
            return {
                "status": "shared",
                "skill_id": skill_id,
                "memory_id": memory_id,
                "name": skill_name,
                "visibility": visibility,
            }
        except Exception as e:
            logger.error("MCP tool_share_skill: %s", e)
            return {"error": "share_failed", "message": str(e)}

    def tool_learn_skill(self, category: str = "", skill_name: str = "",
                         source_agent: str = "unknown",
                         api_key: str | None = None) -> dict[str, Any]:
        """Learn a skill shared by other agents."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=False):
            return {"error": "rate_limited"}

        try:
            # Search for skill packages
            query = f"skill_package {category} {skill_name}".strip()
            results = self.memory.recall(query=query, top_k=5)

            skills = []
            if isinstance(results, dict):
                for mid, mem in results.items():
                    content = mem.get("content", "") if isinstance(mem, dict) else ""
                    if "skill_package" in content:
                        try:
                            skill_data = json.loads(content)
                            if skill_data.get("type") == "skill_package":
                                # Check visibility
                                vis = skill_data.get("visibility", "internal")
                                src = skill_data.get("source_agent", "")
                                if vis == "private" and src != source_agent:
                                    continue
                                skills.append(skill_data)
                        except (json.JSONDecodeError, TypeError):
                            continue
            elif isinstance(results, list):
                for mem in results:
                    content = mem.get("content", "") if isinstance(mem, dict) else ""
                    if "skill_package" in content:
                        try:
                            skill_data = json.loads(content)
                            if skill_data.get("type") == "skill_package":
                                vis = skill_data.get("visibility", "internal")
                                src = skill_data.get("source_agent", "")
                                if vis == "private" and src != source_agent:
                                    continue
                                skills.append(skill_data)
                        except (json.JSONDecodeError, TypeError):
                            continue

            return {
                "status": "ok",
                "skills_found": len(skills),
                "skills": skills,
            }
        except Exception as e:
            logger.error("MCP tool_learn_skill: %s", e)
            return {"error": "learn_failed", "message": str(e)}

    def tool_command(self, command_text: str, confirm: bool = True,
                     api_key: str | None = None) -> dict[str, Any]:
        """Execute a natural language Spirit butler command."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=True):
            return {"error": "rate_limited"}

        if not self.spirit:
            return {"error": "spirit_not_available"}

        try:
            result = self.spirit.execute(command_text, confirm=confirm)
            return {
                "status": "ok" if result.success else "error",
                "intent": getattr(result, 'intent', 'unknown'),
                "action_taken": getattr(result, 'action_taken', ''),
                "output": getattr(result, 'output', ''),
                "error": getattr(result, 'error', None),
            }
        except Exception as e:
            return self._handle_error(e, "tool_command")

    def tool_context_for(self, task_description: str, source_agent: str = "unknown",
                         max_tokens: int = 4000, api_key: str | None = None) -> dict[str, Any]:
        """Assemble a context package for the given task description.

        This is the killer feature for multi-Agent scenarios — context handoff.
        Wraps the existing ContextBuilder if available.
        """
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=False):
            return {"error": "rate_limited"}

        resolved_agent = self._resolve_agent_id(source_agent, api_key)

        try:
            # Try using existing ContextBuilder
            if hasattr(self.memory, 'context_builder') and self.memory.context_builder:
                context = self.memory.context_builder.build(
                    query=task_description,
                    max_tokens=max_tokens,
                )
                return {
                    "status": "ok",
                    "task": task_description,
                    "context": context,
                    "source_agent": resolved_agent,
                }
            # Fallback: use recall as context source
            results = self.memory.recall(query=task_description, top_k=10)
            memories = []
            if isinstance(results, dict):
                for mid, mem in results.items():
                    if isinstance(mem, dict):
                        memories.append({
                            "memory_id": mid,
                            "content": mem.get("content", ""),
                            "importance": mem.get("importance", "medium"),
                        })
            elif isinstance(results, list):
                for mem in results:
                    if isinstance(mem, dict):
                        memories.append(mem)

            # Build context from memories
            context_parts = []
            token_count = 0
            for mem in memories:
                content = mem.get("content", "")
                est_tokens = len(content) // 4
                if token_count + est_tokens > max_tokens:
                    break
                context_parts.append(content)
                token_count += est_tokens

            return {
                "status": "ok",
                "task": task_description,
                "context": "\n---\n".join(context_parts) if context_parts else "",
                "memories_used": len(context_parts),
                "source_agent": resolved_agent,
            }
        except Exception as e:
            return self._handle_error(e, "tool_context_for")

    def tool_correct(self, memory_id: str, correction: str,
                     source_agent: str = "unknown",
                     api_key: str | None = None) -> dict[str, Any]:
        """Correct a memory. Creates a new version with the correction."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=True):
            return {"error": "rate_limited"}

        try:
            if hasattr(self.memory, 'store') and self.memory.store:
                self.memory.store.update_memory(memory_id, correction)
                return {"status": "corrected", "memory_id": memory_id}
            return {"error": "store_not_available"}
        except Exception as e:
            logger.error("MCP tool_correct: %s", e)
            return {"error": "correct_failed", "message": str(e)}

    def tool_delete(self, memory_id: str, source_agent: str = "unknown",
                    api_key: str | None = None) -> dict[str, Any]:
        """Delete a memory by ID."""
        if not self._check_auth(api_key):
            return {"error": "authentication_failed"}
        if not self._check_rate(is_write=True):
            return {"error": "rate_limited"}

        try:
            if hasattr(self.memory, 'store') and self.memory.store:
                self.memory.store.delete_memory(memory_id)
                return {"status": "deleted", "memory_id": memory_id}
            return {"error": "store_not_available"}
        except Exception as e:
            logger.error("MCP tool_delete: %s", e)
            return {"error": "delete_failed", "message": str(e)}

    # --- MCP SDK registration ---

    def _register_tools(self):
        """Register all tools with the FastMCP server."""
        if not self._mcp_server:
            return

        mcp = self._mcp_server

        @mcp.tool()
        def memory_remember(content: str, source_agent: str = "unknown",
                           sensitivity: str = "normal", scope: str = "personal",
                           importance: str = "medium") -> str:
            """Store a memory. The butler auto-tags source agent and manages dedup."""
            result = self.tool_remember(content, source_agent, sensitivity, scope, importance)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_recall(query: str, source_agent: str = "unknown",
                         top_k: int = 10) -> str:
            """Retrieve memories. Privacy-filtered based on agent identity."""
            result = self.tool_recall(query, source_agent, top_k)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_spirit_check() -> str:
            """Spirit proactive check — butler scans for issues and suggests actions."""
            result = self.tool_spirit_check()
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_get_profile(source_agent: str = "unknown") -> str:
            """Get user profile. Spirit decides how much to reveal."""
            result = self.tool_get_profile(source_agent)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_report(report_type: str = "daily") -> str:
            """Generate memory report (daily/weekly/monthly)."""
            result = self.tool_report(report_type)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_share_skill(skill_name: str, skill_description: str,
                              skill_steps: str, category: str = "general",
                              source_agent: str = "unknown",
                              visibility: str = "internal") -> str:
            """Share a skill with other agents. skill_steps should be JSON array string."""
            try:
                steps = json.loads(skill_steps) if isinstance(skill_steps, str) else skill_steps
            except json.JSONDecodeError:
                steps = [skill_steps]
            result = self.tool_share_skill(skill_name, skill_description, steps, category, source_agent, visibility)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_learn_skill(category: str = "", skill_name: str = "",
                              source_agent: str = "unknown") -> str:
            """Learn a skill shared by other agents."""
            result = self.tool_learn_skill(category, skill_name, source_agent)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_command(command_text: str, confirm: bool = True) -> str:
            """Execute a natural language Spirit butler command."""
            result = self.tool_command(command_text, confirm)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_context_for(task_description: str, source_agent: str = "unknown",
                              max_tokens: int = 4000) -> str:
            """Assemble a context package for a task. The killer feature for multi-Agent context handoff."""
            result = self.tool_context_for(task_description, source_agent, max_tokens)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_correct(memory_id: str, correction: str,
                          source_agent: str = "unknown") -> str:
            """Correct a memory. Creates a new version with the correction."""
            result = self.tool_correct(memory_id, correction, source_agent)
            return json.dumps(result, ensure_ascii=False)

        @mcp.tool()
        def memory_delete(memory_id: str, source_agent: str = "unknown") -> str:
            """Delete a memory by ID."""
            result = self.tool_delete(memory_id, source_agent)
            return json.dumps(result, ensure_ascii=False)

    # --- Server lifecycle ---

    def run(self, transport: str | None = None):
        """Start the MCP server.

        Args:
            transport: Override config transport. "stdio" or "streamable-http"
        """
        transport = transport or self.config.transport

        if _MCP_AVAILABLE and self._mcp_server:
            logger.info("Starting MCP server (SDK mode, transport=%s)", transport)
            if transport == "streamable-http":
                self._mcp_server.run(transport="streamable-http")
            else:
                self._mcp_server.run()
        else:
            logger.info("MCP SDK not available, starting HTTP fallback server")
            self._run_http_fallback()

    def _run_http_fallback(self):
        """Run a simple HTTP server as fallback when MCP SDK is not installed."""
        import http.server
        import threading

        class _MCPHTTPHandler(http.server.BaseHTTPRequestHandler):
            server_ref = None  # Set before serving

            def do_POST(self):
                if self.path != "/mcp":
                    self.send_error(404)
                    return

                content_length = int(self.headers.get("Content-Length", 0))
                body = self.rfile.read(content_length).decode("utf-8")

                try:
                    request = json.loads(body)
                except json.JSONDecodeError:
                    self._send_json({"error": "invalid_json"}, 400)
                    return

                tool_name = request.get("tool", "")
                params = request.get("params", {})
                api_key = request.get("api_key") or self.headers.get("X-API-Key")

                tool_map = {
                    "memory.remember": self.server_ref.tool_remember,
                    "memory.recall": self.server_ref.tool_recall,
                    "memory.context_for": self.server_ref.tool_context_for,
                    "memory.correct": self.server_ref.tool_correct,
                    "memory.delete": self.server_ref.tool_delete,
                    "memory.spirit_check": self.server_ref.tool_spirit_check,
                    "memory.get_profile": self.server_ref.tool_get_profile,
                    "memory.report": self.server_ref.tool_report,
                    "memory.share_skill": self.server_ref.tool_share_skill,
                    "memory.learn_skill": self.server_ref.tool_learn_skill,
                    "memory.command": self.server_ref.tool_command,
                }

                if tool_name not in tool_map:
                    self._send_json({"error": f"unknown_tool: {tool_name}"}, 404)
                    return

                if tool_name not in self.server_ref.config.enabled_tools:
                    self._send_json({"error": f"tool_disabled: {tool_name}"}, 403)
                    return

                params["api_key"] = api_key
                try:
                    result = tool_map[tool_name](**params)
                except TypeError as e:
                    self._send_json({"error": f"invalid_params: {e}"}, 400)
                    return

                self._send_json(result)

            def do_GET(self):
                if self.path == "/health":
                    self._send_json({"status": "ok", "version": self.server_ref.config.server_version, "mode": "http_fallback"})
                elif self.path == "/tools":
                    tools = [
                        {"name": t, "description": f"Agent Memory MCP tool: {t}"}
                        for t in self.server_ref.config.enabled_tools
                    ]
                    self._send_json({"tools": tools})
                else:
                    self.send_error(404)

            def _send_json(self, data: dict, status: int = 200):
                self.send_response(status)
                self.send_header("Content-Type", "application/json")
                self.send_header("Access-Control-Allow-Origin", "*")
                self.end_headers()
                self.wfile.write(json.dumps(data, ensure_ascii=False).encode("utf-8"))

            def log_message(self, format, *args):
                logger.debug("MCP HTTP: %s", format % args)

        handler = _MCPHTTPHandler
        handler.server_ref = self

        server = http.server.HTTPServer(
            (self.config.host, self.config.port), handler
        )
        logger.info("MCP HTTP fallback server on %s:%d", self.config.host, self.config.port)
        try:
            server.serve_forever()
        except KeyboardInterrupt:
            server.shutdown()


def create_mcp_server(memory=None, spirit=None, agent_manager=None, config: MCPConfig | None = None) -> AgentMemoryMCPServer:
    """Factory function to create an MCP server from an AgentMemory instance."""
    if memory is None:
        # Try to create a default AgentMemory
        try:
            from .memory_system import AgentMemory
            memory = AgentMemory()
        except Exception as e:
            logger.error("Cannot create default AgentMemory: %s", e)
            raise

    if spirit is None and hasattr(memory, 'spirit'):
        spirit = memory.spirit

    return AgentMemoryMCPServer(memory=memory, spirit=spirit, agent_manager=agent_manager, config=config)
