"""Mem0 memory tools — expose semantic memory to the agent."""

from __future__ import annotations

import json
from typing import Any, TYPE_CHECKING

from loguru import logger

from nanobot.agent.tools.base import Tool

if TYPE_CHECKING:
    from nanobot.agent.memory_mem0 import Mem0MemoryStore


class Mem0ToolContext:
    """Shared mutable state injected into every mem0 tool."""

    def __init__(self, store: Mem0MemoryStore, consolidate_fn):
        self.store = store
        self.consolidate_fn = consolidate_fn  # async (session, archive_all) -> None
        self.user_id: str = "unknown"
        self.session = None

    def set_context(self, channel: str, chat_id: str, session=None):
        self.user_id = f"{channel}_{chat_id}"
        self.session = session


class MemorySearchTool(Tool):
    """Search memories semantically."""

    name = "memory_search"
    description = (
        "Search your long-term memory for facts relevant to a query. "
        "Returns the most relevant memories ranked by similarity."
    )
    parameters = {
        "type": "object",
        "properties": {
            "query": {
                "type": "string",
                "description": "Natural-language search query",
            },
            "limit": {
                "type": "integer",
                "description": "Max results to return (default 5)",
                "minimum": 1,
                "maximum": 20,
            },
        },
        "required": ["query"],
    }

    def __init__(self, ctx: Mem0ToolContext):
        self._ctx = ctx

    async def execute(self, query: str, limit: int = 5, **kw: Any) -> str:
        results = self._ctx.store.search_memories(
            query=query,
            user_id=self._ctx.user_id,
            limit=limit,
        )
        if not results:
            return "No memories found."
        lines = []
        for i, mem in enumerate(results, 1):
            text = mem.get("memory", "")
            score = mem.get("score")
            mid = mem.get("id", "")
            score_str = f"  (score: {score:.2f})" if score else ""
            lines.append(f"{i}. [{mid}] {text}{score_str}")
        return "\n".join(lines)


class MemoryListTool(Tool):
    """List all memories for the current user."""

    name = "memory_list"
    description = (
        "List ALL stored memories for the current user. "
        "Use memory_search for targeted lookup; use this to browse everything."
    )
    parameters = {
        "type": "object",
        "properties": {},
    }

    def __init__(self, ctx: Mem0ToolContext):
        self._ctx = ctx

    async def execute(self, **kw: Any) -> str:
        memories = self._ctx.store.get_all_memories(self._ctx.user_id)
        if not memories:
            return "No memories stored."
        lines = []
        for i, mem in enumerate(memories, 1):
            text = mem.get("memory", "")
            mid = mem.get("id", "")
            lines.append(f"{i}. [{mid}] {text}")
        return f"{len(memories)} memories:\n" + "\n".join(lines)


class MemoryAddTool(Tool):
    """Add a fact to long-term memory."""

    name = "memory_add"
    description = (
        "Store a new fact or piece of information in long-term memory. "
        "The content will be processed by the extraction LLM and stored as one or more facts."
    )
    parameters = {
        "type": "object",
        "properties": {
            "content": {
                "type": "string",
                "description": "The fact or information to remember",
            },
        },
        "required": ["content"],
    }

    def __init__(self, ctx: Mem0ToolContext):
        self._ctx = ctx

    async def execute(self, content: str, **kw: Any) -> str:
        try:
            result = self._ctx.store.memory.add(
                [{"role": "user", "content": content}],
                user_id=self._ctx.user_id,
            )
            facts_count = len(result.get("results", [])) if result else 0
            return f"Added to memory. {facts_count} fact(s) extracted."
        except Exception as e:
            logger.error(f"memory_add failed: {e}")
            return f"Error adding memory: {e}"


class MemoryUpdateTool(Tool):
    """Update an existing memory by ID."""

    name = "memory_update"
    description = (
        "Update the content of an existing memory. "
        "Use memory_list or memory_search first to find the memory ID."
    )
    parameters = {
        "type": "object",
        "properties": {
            "memory_id": {
                "type": "string",
                "description": "The memory ID to update",
            },
            "content": {
                "type": "string",
                "description": "The new content for this memory",
            },
        },
        "required": ["memory_id", "content"],
    }

    def __init__(self, ctx: Mem0ToolContext):
        self._ctx = ctx

    async def execute(self, memory_id: str, content: str, **kw: Any) -> str:
        try:
            self._ctx.store.update_memory(memory_id, content)
            return f"Memory {memory_id} updated."
        except Exception as e:
            logger.error(f"memory_update failed: {e}")
            return f"Error updating memory: {e}"


class MemoryDeleteTool(Tool):
    """Delete a memory by ID."""

    name = "memory_delete"
    description = (
        "Delete a specific memory by its ID. "
        "Use memory_list or memory_search first to find the memory ID."
    )
    parameters = {
        "type": "object",
        "properties": {
            "memory_id": {
                "type": "string",
                "description": "The memory ID to delete",
            },
        },
        "required": ["memory_id"],
    }

    def __init__(self, ctx: Mem0ToolContext):
        self._ctx = ctx

    async def execute(self, memory_id: str, **kw: Any) -> str:
        try:
            self._ctx.store.delete_memory(memory_id)
            return f"Memory {memory_id} deleted."
        except Exception as e:
            logger.error(f"memory_delete failed: {e}")
            return f"Error deleting memory: {e}"


class MemoryConsolidateTool(Tool):
    """Trigger memory consolidation for the current session."""

    name = "memory_consolidate"
    description = (
        "Extract and store facts from the current conversation into long-term memory. "
        "Normally this happens automatically on /new, but you can trigger it manually."
    )
    parameters = {
        "type": "object",
        "properties": {},
    }

    def __init__(self, ctx: Mem0ToolContext):
        self._ctx = ctx

    async def execute(self, **kw: Any) -> str:
        session = self._ctx.session
        if not session:
            return "Error: no active session."
        try:
            await self._ctx.consolidate_fn(session, archive_all=False)
            return "Memory consolidation complete."
        except Exception as e:
            logger.error(f"memory_consolidate failed: {e}")
            return f"Error during consolidation: {e}"
