"""
Memory Bank MCP Server
======================
MCP server providing persistent memory storage for LLM conversations.
Stores and retrieves key-value memories with metadata, tags, and
semantic search capabilities. Backed by a local JSON file store.
"""

import json
import os
import time
import hashlib
from pathlib import Path
from typing import Optional, List, Any, Dict
from datetime import datetime, timezone

from mcp.server.fastmcp import FastMCP

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

MEMORY_DIR = os.environ.get("MEMORY_DIR", "/data/memories")
MAX_MEMORIES = int(os.environ.get("MAX_MEMORIES", "10000"))

# ---------------------------------------------------------------------------
# MCP Server
# ---------------------------------------------------------------------------

mcp = FastMCP("memory_bank_mcp")

# ---------------------------------------------------------------------------
# Storage helpers
# ---------------------------------------------------------------------------


def _ensure_dir():
    """Ensure the memory directory exists."""
    Path(MEMORY_DIR).mkdir(parents=True, exist_ok=True)


def _index_path() -> str:
    return os.path.join(MEMORY_DIR, "_index.json")


def _load_index() -> Dict[str, Any]:
    """Load the memory index."""
    _ensure_dir()
    idx_path = _index_path()
    if os.path.exists(idx_path):
        with open(idx_path, "r") as f:
            return json.load(f)
    return {"memories": {}, "tags": {}}


def _save_index(index: Dict[str, Any]):
    """Save the memory index."""
    _ensure_dir()
    with open(_index_path(), "w") as f:
        json.dump(index, f, indent=2)


def _memory_path(memory_id: str) -> str:
    return os.path.join(MEMORY_DIR, f"{memory_id}.json")


def _generate_id(content: str) -> str:
    """Generate a deterministic ID from content."""
    return hashlib.sha256(content.encode()).hexdigest()[:16]


# ---------------------------------------------------------------------------
# Tools
# ---------------------------------------------------------------------------


@mcp.tool()
async def store_memory(
    key: str,
    content: str,
    tags: Optional[List[str]] = None,
    category: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Store a memory with a key, content, optional tags, category, and metadata.
    If a memory with the same key exists, it will be updated.

    Args:
        key: Unique key/name for this memory
        content: The content to remember
        tags: Optional list of tags for categorization
        category: Optional category (e.g., 'user_preference', 'fact', 'context')
        metadata: Optional additional metadata dict
    """
    index = _load_index()
    now = datetime.now(timezone.utc).isoformat()

    memory = {
        "key": key,
        "content": content,
        "tags": tags or [],
        "category": category or "general",
        "metadata": metadata or {},
        "created_at": now,
        "updated_at": now,
        "access_count": 0,
    }

    # Update if exists
    if key in index["memories"]:
        existing = _load_memory_file(key)
        if existing:
            memory["created_at"] = existing.get("created_at", now)
            memory["access_count"] = existing.get("access_count", 0)

    # Save memory file
    _ensure_dir()
    with open(_memory_path(key), "w") as f:
        json.dump(memory, f, indent=2)

    # Update index
    index["memories"][key] = {
        "category": memory["category"],
        "tags": memory["tags"],
        "updated_at": now,
    }

    # Update tag index
    for tag in memory["tags"]:
        if tag not in index["tags"]:
            index["tags"][tag] = []
        if key not in index["tags"][tag]:
            index["tags"][tag].append(key)

    _save_index(index)

    return {"status": "stored", "key": key, "updated_at": now}


def _load_memory_file(key: str) -> Optional[Dict[str, Any]]:
    """Load a memory file by key."""
    path = _memory_path(key)
    if os.path.exists(path):
        with open(path, "r") as f:
            return json.load(f)
    return None


@mcp.tool()
async def recall_memory(key: str) -> Dict[str, Any]:
    """
    Retrieve a specific memory by its key.

    Args:
        key: The key of the memory to retrieve
    """
    memory = _load_memory_file(key)
    if not memory:
        return {"error": f"Memory '{key}' not found"}

    # Update access count
    memory["access_count"] = memory.get("access_count", 0) + 1
    memory["last_accessed"] = datetime.now(timezone.utc).isoformat()
    with open(_memory_path(key), "w") as f:
        json.dump(memory, f, indent=2)

    return memory


@mcp.tool()
async def search_memories(
    query: Optional[str] = None,
    tags: Optional[List[str]] = None,
    category: Optional[str] = None,
    limit: int = 20,
) -> Dict[str, Any]:
    """
    Search memories by text query, tags, or category.

    Args:
        query: Optional text to search for in memory content and keys
        tags: Optional list of tags to filter by (memories must have ALL tags)
        category: Optional category to filter by
        limit: Maximum results to return (default 20)
    """
    index = _load_index()
    results = []

    # Get candidate keys
    candidate_keys = set(index["memories"].keys())

    # Filter by tags
    if tags:
        for tag in tags:
            tag_keys = set(index["tags"].get(tag, []))
            candidate_keys &= tag_keys

    # Filter by category
    if category:
        candidate_keys = {
            k for k in candidate_keys
            if index["memories"].get(k, {}).get("category") == category
        }

    # Load and search
    for key in candidate_keys:
        memory = _load_memory_file(key)
        if not memory:
            continue

        if query:
            query_lower = query.lower()
            if (
                query_lower not in memory.get("content", "").lower()
                and query_lower not in memory.get("key", "").lower()
                and not any(query_lower in t.lower() for t in memory.get("tags", []))
            ):
                continue

        results.append({
            "key": memory["key"],
            "content": memory["content"][:200] + ("..." if len(memory.get("content", "")) > 200 else ""),
            "category": memory.get("category"),
            "tags": memory.get("tags", []),
            "updated_at": memory.get("updated_at"),
            "access_count": memory.get("access_count", 0),
        })

    # Sort by most recently updated
    results.sort(key=lambda x: x.get("updated_at", ""), reverse=True)

    return {
        "total": len(results),
        "results": results[:limit],
    }


@mcp.tool()
async def delete_memory(key: str) -> Dict[str, Any]:
    """
    Delete a specific memory by key.

    Args:
        key: The key of the memory to delete
    """
    index = _load_index()

    if key not in index["memories"]:
        return {"error": f"Memory '{key}' not found"}

    # Remove from tag index
    mem_info = index["memories"][key]
    for tag in mem_info.get("tags", []):
        if tag in index["tags"] and key in index["tags"][tag]:
            index["tags"][tag].remove(key)
            if not index["tags"][tag]:
                del index["tags"][tag]

    # Remove from index
    del index["memories"][key]
    _save_index(index)

    # Remove file
    path = _memory_path(key)
    if os.path.exists(path):
        os.remove(path)

    return {"status": "deleted", "key": key}


@mcp.tool()
async def list_memories(
    category: Optional[str] = None,
    limit: int = 50,
    offset: int = 0,
) -> Dict[str, Any]:
    """
    List all stored memories with optional category filter.

    Args:
        category: Optional category to filter by
        limit: Maximum results (default 50)
        offset: Skip first N results (default 0)
    """
    index = _load_index()
    memories = []

    for key, info in index["memories"].items():
        if category and info.get("category") != category:
            continue
        memories.append({
            "key": key,
            "category": info.get("category"),
            "tags": info.get("tags", []),
            "updated_at": info.get("updated_at"),
        })

    memories.sort(key=lambda x: x.get("updated_at", ""), reverse=True)
    total = len(memories)

    return {
        "total": total,
        "offset": offset,
        "limit": limit,
        "memories": memories[offset:offset + limit],
    }


@mcp.tool()
async def get_memory_stats() -> Dict[str, Any]:
    """Get statistics about the memory bank."""
    index = _load_index()

    categories = {}
    for info in index["memories"].values():
        cat = info.get("category", "general")
        categories[cat] = categories.get(cat, 0) + 1

    return {
        "total_memories": len(index["memories"]),
        "total_tags": len(index["tags"]),
        "categories": categories,
        "max_memories": MAX_MEMORIES,
    }


@mcp.tool()
async def bulk_store_memories(
    memories: List[Dict[str, Any]],
) -> Dict[str, Any]:
    """
    Store multiple memories at once.

    Args:
        memories: List of memory objects, each with 'key', 'content', and optional 'tags', 'category', 'metadata'
    """
    results = []
    for mem in memories:
        result = await store_memory(
            key=mem["key"],
            content=mem["content"],
            tags=mem.get("tags"),
            category=mem.get("category"),
            metadata=mem.get("metadata"),
        )
        results.append(result)

    return {
        "status": "bulk_stored",
        "count": len(results),
        "results": results,
    }