"""
Memstate AI - LangChain Agent Tools

Provides LangChain ``BaseTool`` implementations that expose the full
Memstate memory API as callable tools for ReAct agents, OpenAI function-calling
agents, and any other LangChain agent framework.

These tools mirror the seven tools provided by the official Memstate MCP server,
making it easy to migrate between MCP and LangChain agent architectures.

Usage::

    from memstate_langchain.tools import create_memstate_tools
    from langchain.agents import create_react_agent, AgentExecutor

    tools = create_memstate_tools(api_key="mst_...", project_id="my-app")
    agent = create_react_agent(llm, tools, prompt)
    executor = AgentExecutor(agent=agent, tools=tools)
"""

from __future__ import annotations

import logging
from typing import Any, Dict, List, Optional, Type

from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

from .memory import MemstateClient, MemstateError

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Input schemas (Pydantic models for LangChain tool argument validation)
# ---------------------------------------------------------------------------


class RememberInput(BaseModel):
    """Input for the memstate_remember tool."""

    content: str = Field(
        ...,
        description=(
            "Markdown or plain text to remember (max 100,000 chars). "
            "The server automatically extracts keypaths and creates structured memories. "
            "Use for task summaries, meeting notes, decisions, or any multi-fact content."
        ),
    )
    source: Optional[str] = Field(
        default="agent",
        description="Source type: agent, readme, docs, meeting, code",
    )
    context: Optional[str] = Field(
        default=None,
        description="Optional hint to guide keypath extraction (e.g. 'authentication setup')",
    )


class SetInput(BaseModel):
    """Input for the memstate_set tool."""

    keypath: str = Field(
        ...,
        description=(
            "Hierarchical dot-path for the value (e.g. 'config.database.port', "
            "'status.deployment', 'version.current'). "
            "Auto-prefixed with 'project.<project_id>.' for you."
        ),
    )
    value: str = Field(
        ...,
        description="Short value to store (max 2,000 chars). For longer content, use memstate_remember.",
    )
    category: Optional[str] = Field(
        default=None,
        description="Category: decision, preference, fact, task, context, requirement, note, code, learning",
    )
    topics: Optional[List[str]] = Field(
        default=None,
        description="Optional additional topics/tags for this memory",
    )


class GetInput(BaseModel):
    """Input for the memstate_get tool."""

    keypath: Optional[str] = Field(
        default=None,
        description=(
            "Subkeypath within the project (e.g. 'database', 'auth.provider'). "
            "If omitted, returns the full project tree."
        ),
    )
    include_content: bool = Field(
        default=False,
        description="If True, include full memory content (not just keypaths). Slower.",
    )
    at_revision: Optional[int] = Field(
        default=None,
        description="Optional revision number for time-travel queries.",
    )


class SearchInput(BaseModel):
    """Input for the memstate_search tool."""

    query: str = Field(..., description="Natural language search query")
    limit: int = Field(default=10, description="Maximum results (default: 10, max: 100)")
    categories: Optional[List[str]] = Field(
        default=None,
        description="Filter by categories (e.g. ['decision', 'fact'])",
    )
    keypath_prefix: Optional[str] = Field(
        default=None,
        description="Filter by keypath prefix (e.g. 'auth')",
    )
    include_superseded: bool = Field(
        default=False,
        description="Include old/superseded memory versions",
    )


class HistoryInput(BaseModel):
    """Input for the memstate_history tool."""

    keypath: Optional[str] = Field(
        default=None,
        description="Get version history for this keypath (e.g. 'config.database.port')",
    )
    memory_id: Optional[str] = Field(
        default=None,
        description="Or get history for a specific memory chain by UUID",
    )


class DeleteInput(BaseModel):
    """Input for the memstate_delete tool."""

    keypath: str = Field(
        ...,
        description=(
            "Keypath to soft-delete (e.g. 'config.old_setting'). "
            "Auto-prefixed with 'project.<project_id>.' for you."
        ),
    )
    recursive: bool = Field(
        default=False,
        description="If True, delete the entire keypath subtree.",
    )


# ---------------------------------------------------------------------------
# Tool implementations
# ---------------------------------------------------------------------------


class MemstateRememberTool(BaseTool):
    """
    Save markdown or text to Memstate AI.

    The server automatically extracts keypaths, detects conflicts, and
    creates structured memories.  This is the preferred way to save
    agent summaries, task results, and multi-fact content.
    """

    name: str = "memstate_remember"
    description: str = (
        "Save markdown, task summaries, or any text to Memstate AI memory. "
        "The server extracts keypaths and creates structured memories automatically. "
        "USE THIS WHEN: saving task summaries, meeting notes, docs, or text with multiple facts. "
        "NOT FOR: setting one specific keypath to a short value (use memstate_set instead). "
        "Content limit: 100,000 chars."
    )
    args_schema: Type[BaseModel] = RememberInput

    client: Any = Field(exclude=True)
    project_id: str

    model_config = {"arbitrary_types_allowed": True}

    def _run(
        self,
        content: str,
        source: Optional[str] = "agent",
        context: Optional[str] = None,
    ) -> str:
        try:
            resp = self.client.ingest(
                project_id=self.project_id,
                content=content,
                source=source or "agent",
                context=context,
            )
            status = resp.get("status", "unknown")
            job_id = resp.get("job_id") or resp.get("ingestion_id", "")
            memories_created = len(resp.get("memories_created", []))
            keypaths_created = resp.get("keypaths_created", [])
            msg = f"Memory saved (status={status}"
            if job_id:
                msg += f", job_id={job_id}"
            if memories_created:
                msg += f", memories_created={memories_created}"
            if keypaths_created:
                msg += f", keypaths={keypaths_created}"
            msg += ")"
            return msg
        except MemstateError as exc:
            return f"Error saving memory: {exc}"


class MemstateSetTool(BaseTool):
    """
    Set one keypath to a short value in Memstate AI.

    Best for simple key=value facts like config, status, or version numbers.
    For rich content or multi-fact summaries, use memstate_remember instead.
    """

    name: str = "memstate_set"
    description: str = (
        "Set one keypath to a short value in Memstate AI. "
        "USE THIS WHEN: storing ONE specific value at a known keypath "
        "(e.g. a port number, status flag, version string). "
        "NOT FOR: task summaries or text with multiple facts (use memstate_remember). "
        "Value limit: 2,000 chars. Keypath is auto-prefixed with 'project.<project_id>.'."
    )
    args_schema: Type[BaseModel] = SetInput

    client: Any = Field(exclude=True)
    project_id: str

    model_config = {"arbitrary_types_allowed": True}

    def _run(
        self,
        keypath: str,
        value: str,
        category: Optional[str] = None,
        topics: Optional[List[str]] = None,
    ) -> str:
        if len(value) > 2000:
            return (
                f"Error: value too long ({len(value)} chars, max 2000). "
                "Use memstate_remember for longer content."
            )
        full_keypath = MemstateClient.build_keypath(self.project_id, keypath)
        try:
            resp = self.client.remember(
                content=value,
                keypath=full_keypath,
                project_id=self.project_id,
                category=category,
                topics=topics,
                source="langchain",
            )
            action = resp.get("action", "stored")
            memory_id = resp.get("memory_id", "")
            version = resp.get("version", 1)
            return (
                f"Memory set at '{full_keypath}' "
                f"(action={action}, memory_id={memory_id}, version={version})"
            )
        except MemstateError as exc:
            return f"Error setting memory: {exc}"


class MemstateGetTool(BaseTool):
    """
    Browse and retrieve memories from Memstate AI.

    Use BEFORE starting tasks to fetch existing project knowledge.
    """

    name: str = "memstate_get"
    description: str = (
        "Browse and retrieve memories from Memstate AI. "
        "USE THIS WHEN: reading what's already stored — list the project tree, "
        "browse a subtree, or get content at a keypath. "
        "NOT FOR: saving content (use memstate_remember or memstate_set). "
        "Examples: memstate_get() → full project tree; "
        "memstate_get(keypath='database') → database subtree."
    )
    args_schema: Type[BaseModel] = GetInput

    client: Any = Field(exclude=True)
    project_id: str

    model_config = {"arbitrary_types_allowed": True}

    def _run(
        self,
        keypath: Optional[str] = None,
        include_content: bool = False,
        at_revision: Optional[int] = None,
    ) -> str:
        full_keypath = MemstateClient.build_keypath(self.project_id, keypath or "")
        try:
            resp = self.client.get_keypaths(
                self.project_id,
                full_keypath,
                recursive=True,
                at_revision=at_revision,
            )
            memories: Dict[str, str] = resp.get("memories", {})
            revision = resp.get("revision_number", 0)
            total = resp.get("total_count", len(memories))

            if not memories:
                return f"No memories found at '{full_keypath}' (revision={revision})"

            lines = [
                f"Project '{self.project_id}' memory tree at '{full_keypath}' "
                f"(total={total}, revision={revision}):"
            ]
            for kp, summary in sorted(memories.items()):
                lines.append(f"  {kp}: {summary}")

            if include_content and memories:
                lines.append("\n--- Full content ---")
                browse_resp = self.client.browse(
                    keypath_prefix=full_keypath,
                    project_id=self.project_id,
                    limit=50,
                )
                for mem in browse_resp.get("memories", []):
                    mem_id = mem.get("id", "")
                    mem_kp = mem.get("keypath", "")
                    if mem_id:
                        try:
                            full = self.client.get_memory(mem_id)
                            content = full.get("content", "")
                            lines.append(f"\n[{mem_kp}]\n{content}")
                        except MemstateError:
                            pass

            return "\n".join(lines)
        except MemstateError as exc:
            return f"Error retrieving memories: {exc}"


class MemstateSearchTool(BaseTool):
    """
    Semantic search across Memstate AI memories.

    Use when you need to find relevant memories by topic or meaning,
    not by exact keypath.
    """

    name: str = "memstate_search"
    description: str = (
        "Find memories by meaning (semantic search) in Memstate AI. "
        "USE THIS WHEN: finding relevant memories by topic or meaning, "
        "not by exact keypath. "
        "NOT FOR: saving content (use memstate_remember or memstate_set). "
        "Returns summaries with similarity scores. "
        "Example: memstate_search(query='how is authentication configured')"
    )
    args_schema: Type[BaseModel] = SearchInput

    client: Any = Field(exclude=True)
    project_id: str

    model_config = {"arbitrary_types_allowed": True}

    def _run(
        self,
        query: str,
        limit: int = 10,
        categories: Optional[List[str]] = None,
        keypath_prefix: Optional[str] = None,
        include_superseded: bool = False,
    ) -> str:
        # Build full keypath prefix if a relative prefix was provided
        if keypath_prefix:
            keypath_prefix = MemstateClient.build_keypath(
                self.project_id, keypath_prefix
            )
        else:
            keypath_prefix = MemstateClient.build_keypath(self.project_id)

        try:
            resp = self.client.search(
                query=query,
                project_id=self.project_id,
                limit=limit,
                categories=categories,
                keypath_prefix=keypath_prefix,
                include_superseded=include_superseded,
            )
            results = resp.get("results", [])
            total = resp.get("total_found", len(results))

            if not results:
                return f"No memories found matching '{query}'"

            lines = [f"Search results for '{query}' (total={total}):"]
            for r in results:
                kp = r.get("keypath", "")
                summary = r.get("summary", r.get("content", ""))
                score = r.get("score", 0.0)
                mem_id = r.get("id", "")
                lines.append(
                    f"  [{kp}] (score={score:.3f}, id={mem_id})\n    {summary}"
                )

            return "\n".join(lines)
        except MemstateError as exc:
            return f"Error searching memories: {exc}"


class MemstateHistoryTool(BaseTool):
    """
    View version history for a keypath or memory chain in Memstate AI.

    Use when debugging past state, auditing changes, or recovering a previous value.
    """

    name: str = "memstate_history"
    description: str = (
        "View version history for a keypath or memory chain in Memstate AI. "
        "USE THIS WHEN: debugging past state, auditing changes, or recovering a previous value. "
        "NOT FOR: reading current content (use memstate_get) or saving (use memstate_remember). "
        "Returns all versions with timestamps; latest is marked is_latest=true."
    )
    args_schema: Type[BaseModel] = HistoryInput

    client: Any = Field(exclude=True)
    project_id: str

    model_config = {"arbitrary_types_allowed": True}

    def _run(
        self,
        keypath: Optional[str] = None,
        memory_id: Optional[str] = None,
    ) -> str:
        if not keypath and not memory_id:
            return "Error: Either keypath or memory_id is required."

        full_keypath = None
        if keypath:
            full_keypath = MemstateClient.build_keypath(self.project_id, keypath)

        try:
            resp = self.client.history(
                memory_id=memory_id,
                keypath=full_keypath,
                project_id=self.project_id,
            )
            versions = resp.get("versions", [])
            total = resp.get("total_versions", len(versions))

            if not versions:
                target = full_keypath or memory_id
                return f"No history found for '{target}'"

            target = full_keypath or memory_id
            lines = [f"Version history for '{target}' (total={total} versions):"]
            for v in versions:
                ver_num = v.get("version", "?")
                created = v.get("created_at", "")
                is_latest = v.get("is_latest", False)
                is_deleted = v.get("is_deleted", False)
                summary = v.get("summary", v.get("content", ""))[:100]
                status = "LATEST" if is_latest else ("DELETED" if is_deleted else "")
                lines.append(
                    f"  v{ver_num} [{created}] {status}\n    {summary}"
                )

            return "\n".join(lines)
        except MemstateError as exc:
            return f"Error retrieving history: {exc}"


class MemstateDeleteTool(BaseTool):
    """
    Soft-delete a memory by keypath in Memstate AI.

    Creates a tombstone version preserving full history.
    """

    name: str = "memstate_delete"
    description: str = (
        "Soft-delete a memory by keypath in Memstate AI. "
        "Creates a tombstone version preserving full history. "
        "USE THIS WHEN: removing outdated or incorrect memories, "
        "cleaning up a keypath subtree, or marking memories as no longer relevant. "
        "NOT FOR: updating content (use memstate_set or memstate_remember to overwrite). "
        "Keypath is auto-prefixed with 'project.<project_id>.'."
    )
    args_schema: Type[BaseModel] = DeleteInput

    client: Any = Field(exclude=True)
    project_id: str

    model_config = {"arbitrary_types_allowed": True}

    def _run(self, keypath: str, recursive: bool = False) -> str:
        full_keypath = MemstateClient.build_keypath(self.project_id, keypath)
        try:
            resp = self.client.delete_memory(
                keypath=full_keypath,
                project_id=self.project_id,
                recursive=recursive,
            )
            deleted_count = resp.get("deleted_count", 0)
            deleted_keypaths = resp.get("deleted_keypaths", [])
            return (
                f"Deleted {deleted_count} memories at '{full_keypath}' "
                f"(keypaths={deleted_keypaths})"
            )
        except MemstateError as exc:
            return f"Error deleting memory: {exc}"


# ---------------------------------------------------------------------------
# Factory function
# ---------------------------------------------------------------------------


def create_memstate_tools(
    api_key: str,
    project_id: str,
    *,
    base_url: str = "https://api.memstate.ai",
    timeout: int = 60,
    include_tools: Optional[List[str]] = None,
) -> List[BaseTool]:
    """
    Create a list of LangChain tools for interacting with Memstate AI.

    This is the primary entry point for adding Memstate memory tools to a
    LangChain agent.  Returns all seven tools by default, matching the
    official Memstate MCP server's tool set.

    Args:
        api_key: Your Memstate API key (get one at https://memstate.ai/dashboard).
        project_id: The Memstate project to use for all tool calls.
        base_url: Memstate API base URL (default: https://api.memstate.ai).
        timeout: HTTP request timeout in seconds.
        include_tools: Optional list of tool names to include. If None, all
            tools are included. Valid names: ``memstate_remember``,
            ``memstate_set``, ``memstate_get``, ``memstate_search``,
            ``memstate_history``, ``memstate_delete``.

    Returns:
        List of :class:`~langchain_core.tools.BaseTool` instances.

    Example::

        from memstate_langchain.tools import create_memstate_tools
        from langchain_openai import ChatOpenAI
        from langchain.agents import create_tool_calling_agent, AgentExecutor
        from langchain_core.prompts import ChatPromptTemplate

        tools = create_memstate_tools(api_key="mst_...", project_id="my-app")

        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful assistant with persistent memory."),
            ("placeholder", "{chat_history}"),
            ("human", "{input}"),
            ("placeholder", "{agent_scratchpad}"),
        ])

        agent = create_tool_calling_agent(ChatOpenAI(model="gpt-4o"), tools, prompt)
        executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
    """
    client = MemstateClient(api_key=api_key, base_url=base_url, timeout=timeout)

    all_tools: List[BaseTool] = [
        MemstateRememberTool(client=client, project_id=project_id),
        MemstateSetTool(client=client, project_id=project_id),
        MemstateGetTool(client=client, project_id=project_id),
        MemstateSearchTool(client=client, project_id=project_id),
        MemstateHistoryTool(client=client, project_id=project_id),
        MemstateDeleteTool(client=client, project_id=project_id),
    ]

    if include_tools is None:
        return all_tools

    name_set = set(include_tools)
    return [t for t in all_tools if t.name in name_set]
