from __future__ import annotations

import argparse
import json
import math
import os
import re
import sqlite3
from datetime import datetime, timedelta, timezone
from pathlib import Path

from mcp.server import MCPServer


DB_PATH = Path(os.getenv("AGENT_MEMORY_DB", "memory.db")).resolve()
MAX_CONTENT_LENGTH = 2_000
MAX_RECALL_LIMIT = 20

mcp = MCPServer("agent-memory")


def now_utc() -> datetime:
    return datetime.now(timezone.utc)


def iso_now() -> str:
    return now_utc().isoformat(timespec="seconds")


def tokens(text: str) -> set[str]:
    """保留中文单字、英文单词和数字，用于轻量本地召回。"""
    return set(re.findall(r"[\u4e00-\u9fff]|[a-zA-Z0-9_]+", text.casefold()))


def connect() -> sqlite3.Connection:
    DB_PATH.parent.mkdir(parents=True, exist_ok=True)
    db = sqlite3.connect(DB_PATH)
    db.row_factory = sqlite3.Row
    db.execute("PRAGMA journal_mode=WAL")
    db.execute("PRAGMA foreign_keys=ON")
    db.executescript(
        """
        CREATE TABLE IF NOT EXISTS memories (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            user_id TEXT NOT NULL,
            content TEXT NOT NULL,
            category TEXT NOT NULL DEFAULT 'fact',
            importance INTEGER NOT NULL CHECK (importance BETWEEN 1 AND 5),
            source TEXT NOT NULL DEFAULT 'agent',
            status TEXT NOT NULL DEFAULT 'active'
                CHECK (status IN ('active', 'superseded', 'deleted')),
            supersedes_id INTEGER REFERENCES memories(id),
            created_at TEXT NOT NULL,
            updated_at TEXT NOT NULL,
            expires_at TEXT,
            deleted_at TEXT
        );

        CREATE INDEX IF NOT EXISTS idx_memories_user_status
            ON memories(user_id, status, updated_at);

        CREATE TABLE IF NOT EXISTS memory_events (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            memory_id INTEGER,
            user_id TEXT NOT NULL,
            action TEXT NOT NULL,
            detail TEXT NOT NULL,
            created_at TEXT NOT NULL,
            FOREIGN KEY(memory_id) REFERENCES memories(id)
        );
        """
    )
    return db


def validate_user(user_id: str) -> str:
    user_id = user_id.strip()
    if not user_id or len(user_id) > 100:
        raise ValueError("user_id 必须是 1-100 个字符")
    return user_id


def validate_content(content: str) -> str:
    content = content.strip()
    if not content:
        raise ValueError("content 不能为空")
    if len(content) > MAX_CONTENT_LENGTH:
        raise ValueError(f"content 最多 {MAX_CONTENT_LENGTH} 个字符")
    return content


def audit(
    db: sqlite3.Connection,
    memory_id: int | None,
    user_id: str,
    action: str,
    detail: dict,
) -> None:
    db.execute(
        "INSERT INTO memory_events(memory_id, user_id, action, detail, created_at) VALUES (?, ?, ?, ?, ?)",
        (memory_id, user_id, action, json.dumps(detail, ensure_ascii=False), iso_now()),
    )


def remember_memory(
    user_id: str,
    content: str,
    category: str = "fact",
    importance: int = 3,
    expires_in_days: int | None = None,
    source: str = "agent",
    supersedes_id: int | None = None,
) -> dict:
    user_id = validate_user(user_id)
    content = validate_content(content)
    category = category.strip() or "fact"
    source = source.strip() or "agent"
    if not 1 <= importance <= 5:
        raise ValueError("importance 必须是 1-5")
    if expires_in_days is not None and not 1 <= expires_in_days <= 3650:
        raise ValueError("expires_in_days 必须是 1-3650")

    created_at = iso_now()
    expires_at = None
    if expires_in_days is not None:
        expires_at = (now_utc() + timedelta(days=expires_in_days)).isoformat(timespec="seconds")

    with connect() as db:
        duplicate = db.execute(
            """
            SELECT id FROM memories
            WHERE user_id = ? AND content = ? AND status = 'active'
            ORDER BY id DESC LIMIT 1
            """,
            (user_id, content),
        ).fetchone()
        if duplicate:
            return {"id": duplicate["id"], "status": "already_exists"}

        if supersedes_id is not None:
            old = db.execute(
                "SELECT id FROM memories WHERE id = ? AND user_id = ? AND status = 'active'",
                (supersedes_id, user_id),
            ).fetchone()
            if not old:
                raise ValueError("supersedes_id 不属于当前用户，或记忆已不是 active")

        cursor = db.execute(
            """
            INSERT INTO memories(
                user_id, content, category, importance, source, supersedes_id,
                created_at, updated_at, expires_at
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            (
                user_id,
                content,
                category,
                importance,
                source,
                supersedes_id,
                created_at,
                created_at,
                expires_at,
            ),
        )
        memory_id = cursor.lastrowid
        if supersedes_id is not None:
            db.execute(
                "UPDATE memories SET status = 'superseded', updated_at = ? WHERE id = ?",
                (created_at, supersedes_id),
            )
        audit(
            db,
            memory_id,
            user_id,
            "remember",
            {"category": category, "importance": importance, "supersedes_id": supersedes_id},
        )
        if supersedes_id is not None:
            audit(db, supersedes_id, user_id, "superseded", {"by": memory_id})
        return {"id": memory_id, "status": "created", "expires_at": expires_at}


def recall_memories(user_id: str, query: str, limit: int = 5) -> list[dict]:
    user_id = validate_user(user_id)
    query = query.strip()
    if not query:
        raise ValueError("query 不能为空")
    limit = max(1, min(limit, MAX_RECALL_LIMIT))
    query_tokens = tokens(query)
    now = now_utc()
    with connect() as db:
        rows = db.execute(
            "SELECT * FROM memories WHERE user_id = ? AND status = 'active'",
            (user_id,),
        ).fetchall()

    candidates: list[dict] = []
    for row in rows:
        if row["expires_at"] and datetime.fromisoformat(row["expires_at"]) <= now:
            continue
        memory_tokens = tokens(row["content"])
        overlap = len(query_tokens & memory_tokens) / max(len(query_tokens), 1)
        if overlap == 0:
            continue
        age_days = max((now - datetime.fromisoformat(row["updated_at"])).total_seconds() / 86400, 0)
        recency = math.exp(-age_days / 30)
        score = 0.65 * overlap + 0.20 * (row["importance"] / 5) + 0.15 * recency
        candidates.append(
            {
                "id": row["id"],
                "content": row["content"],
                "category": row["category"],
                "importance": row["importance"],
                "score": round(score, 4),
                "created_at": row["created_at"],
                "expires_at": row["expires_at"],
                "source": row["source"],
            }
        )
    return sorted(candidates, key=lambda item: item["score"], reverse=True)[:limit]


def forget_memory(user_id: str, memory_id: int, reason: str = "user_request") -> dict:
    user_id = validate_user(user_id)
    reason = reason.strip()[:200] or "user_request"
    with connect() as db:
        row = db.execute(
            "SELECT id FROM memories WHERE id = ? AND user_id = ? AND status = 'active'",
            (memory_id, user_id),
        ).fetchone()
        if not row:
            raise ValueError("找不到当前用户的 active 记忆")
        deleted_at = iso_now()
        db.execute(
            "UPDATE memories SET status = 'deleted', deleted_at = ?, updated_at = ? WHERE id = ?",
            (deleted_at, deleted_at, memory_id),
        )
        audit(db, memory_id, user_id, "forget", {"reason": reason})
    return {"id": memory_id, "status": "deleted", "deleted_at": deleted_at}


@mcp.tool()
def remember(
    user_id: str,
    content: str,
    category: str = "fact",
    importance: int = 3,
    expires_in_days: int | None = None,
    source: str = "agent",
    supersedes_id: int | None = None,
) -> dict:
    """保存一条记忆；新记忆可以用 supersedes_id 替代旧记忆。"""
    return remember_memory(user_id, content, category, importance, expires_in_days, source, supersedes_id)


@mcp.tool()
def recall(user_id: str, query: str, limit: int = 5) -> list[dict]:
    """召回当前用户的 active 记忆，自动过滤过期和已删除内容。"""
    return recall_memories(user_id, query, limit)


@mcp.tool()
def forget(user_id: str, memory_id: int, reason: str = "user_request") -> dict:
    """软删除一条记忆；删除动作会写入审计表。"""
    return forget_memory(user_id, memory_id, reason)


def cli() -> None:
    parser = argparse.ArgumentParser(description="local Agent memory")
    sub = parser.add_subparsers(dest="command", required=True)

    remember_parser = sub.add_parser("remember")
    remember_parser.add_argument("--user", required=True, dest="user_id")
    remember_parser.add_argument("--content", required=True)
    remember_parser.add_argument("--category", default="fact")
    remember_parser.add_argument("--importance", type=int, default=3)
    remember_parser.add_argument("--expires-in-days", type=int)
    remember_parser.add_argument("--source", default="cli")
    remember_parser.add_argument("--supersedes-id", type=int)

    recall_parser = sub.add_parser("recall")
    recall_parser.add_argument("--user", required=True, dest="user_id")
    recall_parser.add_argument("--query", required=True)
    recall_parser.add_argument("--limit", type=int, default=5)

    forget_parser = sub.add_parser("forget")
    forget_parser.add_argument("--user", required=True, dest="user_id")
    forget_parser.add_argument("--memory-id", type=int, required=True)
    forget_parser.add_argument("--reason", default="user_request")

    args = parser.parse_args()
    if args.command == "remember":
        result = remember_memory(
            args.user_id,
            args.content,
            args.category,
            args.importance,
            args.expires_in_days,
            args.source,
            args.supersedes_id,
        )
    elif args.command == "recall":
        result = recall_memories(args.user_id, args.query, args.limit)
    else:
        result = forget_memory(args.user_id, args.memory_id, args.reason)
    print(json.dumps(result, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    cli()

import asyncio
import tempfile
from pathlib import Path

from mcp import Client

import memory


async def main() -> None:
    with tempfile.TemporaryDirectory() as directory:
        memory.DB_PATH = Path(directory) / "test.db"
        async with Client(memory.mcp) as client:
            created = await client.call_tool(
                "remember",
                {"user_id": "test", "content": "喜欢简洁的 Python 示例", "category": "preference"},
            )
            assert created.structured_content["id"] == 1

            found = await client.call_tool(
                "recall", {"user_id": "test", "query": "Python 示例"}
            )
            assert found.structured_content[0]["content"] == "喜欢简洁的 Python 示例"

            await client.call_tool("forget", {"user_id": "test", "memory_id": 1})
            empty = await client.call_tool(
                "recall", {"user_id": "test", "query": "Python 示例"}
            )
            assert empty.structured_content == []


asyncio.run(main())