from __future__ import annotations

import json
from typing import TYPE_CHECKING

from niannian.mcp.mcp_space import resolve_mcp_memory_space
from niannian.memory.space import MemorySpaceError, require_mcp_memory_space

if TYPE_CHECKING:
    from mcp.server.fastmcp import FastMCP

    from niannian.mcp.service import NianNianService


def register_recall_tools(mcp: FastMCP, svc: NianNianService) -> None:
    @mcp.tool()
    def niannian_assist(task: str, project: str, context: str = "") -> str:
        """念念：碎片化思考辅助（= recall assist_default）。project=记忆空间（必填）。"""
        space = resolve_mcp_memory_space(svc, project)
        return json.dumps(svc.assist(task, context=context, project=space), ensure_ascii=False)

    @mcp.tool()
    def niannian_query(question: str, project: str) -> str:
        """念念：基于记忆回答问题。project=记忆空间（必填）。"""
        space = resolve_mcp_memory_space(svc, project)
        return json.dumps(svc.query(question, project=space), ensure_ascii=False)

    @mcp.tool()
    def niannian_record(fragment_json: str) -> str:
        """MRRR 记录 Fragment/2。fragment 内 project=记忆空间（必填）。"""
        data = json.loads(fragment_json)
        try:
            space = require_mcp_memory_space(data.get("project"))
        except MemorySpaceError as exc:
            raise ValueError(str(exc)) from exc
        home = svc.home
        if home is not None:
            from niannian.memory.space import ensure_memory_space

            ensure_memory_space(home, space)
        data["project"] = space
        return json.dumps(svc.record(fragment=data), ensure_ascii=False)

    @mcp.tool()
    def niannian_recall(
        task: str,
        project: str,
        policy: str = "assist_default",
        scope_json: str = "{}",
    ) -> str:
        """MRRR 回忆。project=记忆空间（必填）。loop 续写 policy=loop_continue。"""
        scope = json.loads(scope_json) if scope_json else {}
        space = resolve_mcp_memory_space(svc, project)
        return json.dumps(
            svc.recall(task, scope=scope, policy=policy, project=space),
            ensure_ascii=False,
        )

    @mcp.tool()
    def niannian_reflect(project: str, scope_json: str = "{}") -> str:
        """MRRR 反思 — thread 摘要 + dream。project=记忆空间（必填）。"""
        scope = json.loads(scope_json) if scope_json else {}
        space = resolve_mcp_memory_space(svc, project)
        return json.dumps(svc.reflect(scope=scope, project=space), ensure_ascii=False)
