"""
Memstate AI - LangChain Memory Integration

Provides LangChain-compatible memory classes backed by the Memstate AI REST API.
Memstate uses a hierarchical keypath system for structured, versioned memory storage
with semantic search — a direct competitor to Mem0 with superior conflict detection
and version history.

This module targets LangChain 0.3+ (langchain-core >= 0.3.0), which removed the
deprecated ``BaseMemory`` class in favor of the modern LCEL Runnable pattern.

The three integration patterns provided are:

1. ``MemstateChatMessageHistory`` — drop-in ``BaseChatMessageHistory`` for
   ``RunnableWithMessageHistory``, persisting the full message list.

2. ``MemstateSemanticMemory`` — a ``RunnableLambda``-based semantic memory
   injector that retrieves relevant memories before each chain call and
   persists the conversation turn afterward.  Designed to wrap any LCEL chain.

3. ``MemstateMemoryRunnable`` — a convenience factory that returns a
   ``RunnableWithMessageHistory``-wrapped chain augmented with semantic memory
   injection, providing the closest equivalent to the old ``ConversationChain``
   with ``MemstateMemory``.

API Reference: https://memstate.ai/docs
"""

from __future__ import annotations

import json
import logging
import time
from typing import Any, Callable, Dict, List, Optional, Sequence, Union

import requests
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.messages import (
    AIMessage,
    BaseMessage,
    HumanMessage,
    messages_from_dict,
    messages_to_dict,
)
from langchain_core.runnables import (
    Runnable,
    RunnableLambda,
    RunnablePassthrough,
)
from langchain_core.runnables.history import RunnableWithMessageHistory
from pydantic import BaseModel, Field, model_validator

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

DEFAULT_BASE_URL = "https://api.memstate.ai"
DEFAULT_TIMEOUT = 60  # seconds
DEFAULT_SEARCH_LIMIT = 10
DEFAULT_INGEST_POLL_INTERVAL = 2.0  # seconds between job status polls
DEFAULT_INGEST_MAX_WAIT = 120.0  # max seconds to wait for async ingest


# ---------------------------------------------------------------------------
# Exceptions
# ---------------------------------------------------------------------------


class MemstateError(Exception):
    """Base exception for all Memstate errors."""


class MemstateAPIError(MemstateError):
    """Raised when the Memstate API returns a non-2xx response."""

    def __init__(self, message: str, status_code: Optional[int] = None) -> None:
        super().__init__(message)
        self.status_code = status_code


class MemstateConnectionError(MemstateError):
    """Raised when a network-level error prevents reaching the Memstate API."""


# ---------------------------------------------------------------------------
# Low-level REST client
# ---------------------------------------------------------------------------


class MemstateClient:
    """
    Thin, synchronous HTTP client for the Memstate AI REST API.

    Authentication uses the ``X-API-Key`` header, which is the same mechanism
    used by the official MCP server and the Go backend.  Both
    ``Authorization: Bearer <key>`` and ``X-API-Key: <key>`` are accepted by
    the server; this client uses ``X-API-Key`` to match the MCP convention.

    All public methods raise :class:`MemstateAPIError` on non-2xx responses.
    """

    def __init__(
        self,
        api_key: str,
        base_url: str = DEFAULT_BASE_URL,
        timeout: int = DEFAULT_TIMEOUT,
    ) -> None:
        if not api_key:
            raise ValueError(
                "api_key is required. Get yours at https://memstate.ai/dashboard"
            )
        self.api_key = api_key
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._session = requests.Session()
        self._session.headers.update(
            {
                "X-API-Key": api_key,
                "Content-Type": "application/json",
                "User-Agent": "memstate-langchain/1.0.0",
            }
        )

    # ------------------------------------------------------------------
    # Internal helpers
    # ------------------------------------------------------------------

    def _request(
        self,
        method: str,
        path: str,
        *,
        json_body: Optional[Dict[str, Any]] = None,
        params: Optional[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        url = f"{self.base_url}{path}"
        try:
            resp = self._session.request(
                method,
                url,
                json=json_body,
                params=params,
                timeout=self.timeout,
            )
        except requests.RequestException as exc:
            raise MemstateConnectionError(
                f"Failed to connect to Memstate at {self.base_url}: {exc}"
            ) from exc

        if not resp.ok:
            try:
                detail = resp.json().get("error", resp.text)
            except Exception:
                detail = resp.text
            raise MemstateAPIError(
                f"Memstate API error {resp.status_code}: {detail}",
                status_code=resp.status_code,
            )

        try:
            return resp.json()
        except Exception as exc:
            raise MemstateAPIError(
                f"Failed to parse Memstate response: {exc}"
            ) from exc

    # ------------------------------------------------------------------
    # Memory endpoints
    # ------------------------------------------------------------------

    def remember(
        self,
        content: str,
        keypath: str,
        *,
        project_id: Optional[str] = None,
        category: Optional[str] = None,
        topics: Optional[List[str]] = None,
        source: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        POST /api/v1/memories/remember

        Store a single memory at an explicit keypath.  Best for short,
        structured values (config, status, version numbers).  For rich
        markdown or multi-fact content, use :meth:`ingest` instead.
        """
        body: Dict[str, Any] = {"content": content, "keypath": keypath}
        if project_id is not None:
            body["project_id"] = project_id
        if category:
            body["category"] = category
        if topics:
            body["topics"] = topics
        if source:
            body["source"] = source
        return self._request("POST", "/api/v1/memories/remember", json_body=body)

    def ingest(
        self,
        project_id: str,
        content: str,
        *,
        source: Optional[str] = None,
        context: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        POST /api/v1/ingest

        Server-side intelligence endpoint: submit raw markdown or text and
        let the Memstate LLM extract keypaths, detect conflicts, and create
        structured memories automatically.  This is the preferred write path
        for agent summaries, meeting notes, and multi-fact content.

        Returns immediately with a ``job_id`` if async processing is enabled,
        or with the full result if synchronous.
        """
        body: Dict[str, Any] = {"project_id": project_id, "content": content}
        if source:
            body["source"] = source
        if context:
            body["context"] = context
        return self._request("POST", "/api/v1/ingest", json_body=body)

    def get_job_status(self, job_id: str) -> Dict[str, Any]:
        """GET /jobs/{job_id} — poll async ingest job status."""
        return self._request("GET", f"/jobs/{job_id}")

    def search(
        self,
        query: str,
        *,
        project_id: Optional[str] = None,
        limit: int = DEFAULT_SEARCH_LIMIT,
        categories: Optional[List[str]] = None,
        keypath_prefix: Optional[str] = None,
        include_superseded: bool = False,
    ) -> Dict[str, Any]:
        """
        POST /api/v1/memories/search

        Semantic (vector) search across memories.  Returns ranked results
        with similarity scores.

        The Memstate API returns results in the format::

            {"results": [{"memory": {"id": ..., "keypath": ..., "summary": ...}, "score": 0.9}]}

        This method normalizes each result by merging the nested ``memory``
        object with the top-level ``score``, so callers receive flat dicts::

            {"id": ..., "keypath": ..., "summary": ..., "score": 0.9}
        """
        body: Dict[str, Any] = {"query": query, "limit": limit}
        if project_id is not None:
            body["project_id"] = project_id
        if categories:
            body["categories"] = categories
        if keypath_prefix is not None:
            body["keypath_prefix"] = keypath_prefix
        if include_superseded:
            body["include_superseded"] = True
        raw = self._request("POST", "/api/v1/memories/search", json_body=body)
        # Normalize nested {memory: {...}, score: float} → flat dict with score
        normalized_results = []
        for item in raw.get("results", []):
            if isinstance(item, dict) and "memory" in item:
                flat = dict(item["memory"])
                flat["score"] = item.get("score", 0.0)
                normalized_results.append(flat)
            else:
                normalized_results.append(item)
        raw["results"] = normalized_results
        return raw

    def browse(
        self,
        keypath_prefix: str,
        *,
        project_id: Optional[str] = None,
        limit: int = 100,
    ) -> Dict[str, Any]:
        """POST /api/v1/memories/browse — list memories under a keypath prefix."""
        body: Dict[str, Any] = {"keypath_prefix": keypath_prefix, "limit": limit}
        if project_id is not None:
            body["project_id"] = project_id
        return self._request("POST", "/api/v1/memories/browse", json_body=body)

    def get_memory(self, memory_id: str) -> Dict[str, Any]:
        """GET /api/v1/memories/{id} — fetch a single memory by UUID."""
        return self._request("GET", f"/api/v1/memories/{memory_id}")

    def get_keypaths(
        self,
        project_id: str,
        keypath: str,
        *,
        recursive: bool = True,
        at_revision: Optional[int] = None,
    ) -> Dict[str, Any]:
        """
        POST /api/v1/keypaths

        Get the keypath tree for a project (or subtree).  Supports
        time-travel via ``at_revision``.
        """
        body: Dict[str, Any] = {
            "project_id": project_id,
            "keypath": keypath,
            "recursive": recursive,
        }
        if at_revision is not None:
            body["at_revision"] = at_revision
        return self._request("POST", "/api/v1/keypaths", json_body=body)

    def history(
        self,
        *,
        memory_id: Optional[str] = None,
        keypath: Optional[str] = None,
        project_id: Optional[str] = None,
    ) -> Dict[str, Any]:
        """POST /api/v1/memories/history — version history for a memory or keypath."""
        if not memory_id and not keypath:
            raise ValueError("Either memory_id or keypath is required")
        body: Dict[str, Any] = {}
        if memory_id:
            body["memory_id"] = memory_id
        if keypath:
            body["keypath"] = keypath
        if project_id:
            body["project_id"] = project_id
        return self._request("POST", "/api/v1/memories/history", json_body=body)

    def supersede(
        self,
        memory_id: str,
        content: str,
        *,
        category: Optional[str] = None,
    ) -> Dict[str, Any]:
        """POST /api/v1/memories/supersede — update a memory, creating a new version."""
        body: Dict[str, Any] = {"memory_id": memory_id, "content": content}
        if category:
            body["category"] = category
        return self._request("POST", "/api/v1/memories/supersede", json_body=body)

    def delete_memory(
        self,
        keypath: str,
        project_id: str,
        *,
        recursive: bool = False,
    ) -> Dict[str, Any]:
        """POST /api/v1/memories/delete — soft-delete a keypath (tombstone)."""
        body: Dict[str, Any] = {
            "keypath": keypath,
            "project_id": project_id,
            "recursive": recursive,
        }
        return self._request("POST", "/api/v1/memories/delete", json_body=body)

    def list_projects(self) -> Dict[str, Any]:
        """GET /api/v1/projects — list all projects."""
        return self._request("GET", "/api/v1/projects")

    def create_project(
        self,
        project_id: str,
        *,
        name: Optional[str] = None,
        description: Optional[str] = None,
    ) -> Dict[str, Any]:
        """POST /api/v1/projects — create or upsert a project."""
        body: Dict[str, Any] = {"project_id": project_id}
        if name:
            body["name"] = name
        if description:
            body["description"] = description
        return self._request("POST", "/api/v1/projects", json_body=body)

    def get_tree(self, project_id: Optional[str] = None) -> Dict[str, Any]:
        """GET /api/v1/tree — get the top-level domain tree."""
        params = {}
        if project_id:
            params["project_id"] = project_id
        return self._request("GET", "/api/v1/tree", params=params or None)

    def get_status(self) -> Dict[str, Any]:
        """GET /api/v1/status — system health and memory counts."""
        return self._request("GET", "/api/v1/status")

    # ------------------------------------------------------------------
    # Convenience: build the standard keypath prefix used by the MCP server
    # ------------------------------------------------------------------

    @staticmethod
    def build_keypath(project_id: str, keypath: str = "") -> str:
        """
        Mirror the ``buildKeypath`` function from the Go MCP server.

        Keypaths are auto-prefixed with ``project.<project_id>.`` so that
        memories are namespaced per project, matching the MCP convention.

        The Memstate API requires keypaths to contain only lowercase letters,
        digits, underscores, and dots.  Hyphens are not permitted.  This method
        automatically sanitizes the project_id component by replacing hyphens
        with underscores so that project IDs like ``my-app`` produce valid
        keypaths like ``project.my_app.database.type``.
        """
        if not project_id:
            return keypath
        # Sanitize: replace hyphens with underscores in the keypath segment
        safe_project_id = project_id.replace("-", "_")
        prefix = f"project.{safe_project_id}"
        if not keypath:
            return prefix
        return f"{prefix}.{keypath}"


# ---------------------------------------------------------------------------
# LangChain Chat Message History
# ---------------------------------------------------------------------------


class MemstateChatMessageHistory(BaseChatMessageHistory):
    """
    Persistent chat message history backed by Memstate AI.

    Stores the full conversation as a JSON blob at a deterministic keypath
    (``project.<project_id>.chat_history.<session_id>``), enabling multi-session
    memory that persists across agent restarts.

    This is a drop-in replacement for ``RedisChatMessageHistory``,
    ``DynamoDBChatMessageHistory``, or Mem0's ``MemoryHistory``.

    Note: This class intentionally does NOT inherit from ``pydantic.BaseModel``
    because Pydantic v2 does not support ``@property`` overrides on model fields.
    The ``messages`` property must be dynamic (fetched from the Memstate API on
    each access), which requires a plain Python class.

    Example::

        from memstate_langchain import MemstateChatMessageHistory
        from langchain_core.runnables.history import RunnableWithMessageHistory

        chain_with_history = RunnableWithMessageHistory(
            chain,
            lambda session_id: MemstateChatMessageHistory(
                api_key="mst_...",
                project_id="my-assistant",
                session_id=session_id,
            ),
            input_messages_key="input",
            history_messages_key="history",
        )
    """

    def __init__(
        self,
        api_key: str,
        project_id: str,
        session_id: str,
        *,
        base_url: str = DEFAULT_BASE_URL,
        timeout: int = DEFAULT_TIMEOUT,
    ) -> None:
        """
        Args:
            api_key: Memstate API key.
            project_id: Memstate project ID.
            session_id: Unique session/conversation ID.
            base_url: Memstate API base URL.
            timeout: HTTP timeout in seconds.
        """
        self.api_key = api_key
        self.project_id = project_id
        self.session_id = session_id
        self.base_url = base_url
        self.timeout = timeout
        self._client = MemstateClient(
            api_key=api_key,
            base_url=base_url,
            timeout=timeout,
        )

    @property
    def _keypath(self) -> str:
        """
        Full keypath for this session's chat history.

        The Memstate API requires keypaths to contain only lowercase letters,
        digits, underscores, and dots.  Hyphens in the session_id are
        sanitized to underscores so that session IDs like ``user-123`` produce
        valid keypaths like ``project.my_app.chat_history.user_123``.
        """
        # Sanitize session_id: replace hyphens with underscores
        safe_session_id = self.session_id.replace("-", "_")
        return MemstateClient.build_keypath(
            self.project_id, f"chat_history.{safe_session_id}"
        )

    @property
    def messages(self) -> List[BaseMessage]:
        """
        Retrieve all messages for this session from Memstate.

        The full conversation history is stored as a JSON-serialised list at
        the session's keypath.  The ``browse`` endpoint returns a truncated
        ``summary`` field, so we first browse to find the memory ID, then
        fetch the full memory via ``GET /api/v1/memories/{id}`` to obtain the
        complete ``content`` field.
        """
        assert self._client is not None
        try:
            resp = self._client.browse(
                keypath_prefix=self._keypath,
                project_id=self.project_id,
                limit=1,
            )
            memories_list = resp.get("memories", [])
            if not memories_list:
                return []
            # Find the memory whose keypath exactly matches our session keypath
            memory_id = None
            for mem in memories_list:
                if isinstance(mem, dict) and mem.get("keypath") == self._keypath:
                    memory_id = mem.get("id")
                    break
            if not memory_id:
                return []
            # Fetch the full memory to get the complete content (summary is truncated)
            full_mem = self._client.get_memory(memory_id)
            raw = full_mem.get("content", "")
            if not raw:
                return []
            data = json.loads(raw)
            return messages_from_dict(data)
        except MemstateAPIError as exc:
            if exc.status_code == 404:
                return []
            logger.warning("Failed to load chat history from Memstate: %s", exc)
            return []
        except Exception as exc:
            logger.warning("Failed to parse chat history from Memstate: %s", exc)
            return []

    def add_message(self, message: BaseMessage) -> None:
        """Append a message and persist the updated history to Memstate."""
        assert self._client is not None
        current = self.messages
        current.append(message)
        serialized = json.dumps(messages_to_dict(current))
        try:
            self._client.remember(
                content=serialized,
                keypath=self._keypath,
                project_id=self.project_id,
                category="context",
                source="langchain",
            )
        except MemstateError as exc:
            logger.error("Failed to persist chat message to Memstate: %s", exc)
            raise

    def add_messages(self, messages: Sequence[BaseMessage]) -> None:
        """Append multiple messages in a single write (more efficient than add_message)."""
        assert self._client is not None
        current = self.messages
        current.extend(messages)
        serialized = json.dumps(messages_to_dict(current))
        try:
            self._client.remember(
                content=serialized,
                keypath=self._keypath,
                project_id=self.project_id,
                category="context",
                source="langchain",
            )
        except MemstateError as exc:
            logger.error("Failed to persist chat messages to Memstate: %s", exc)
            raise

    def clear(self) -> None:
        """Delete the chat history for this session from Memstate."""
        assert self._client is not None
        try:
            self._client.delete_memory(
                keypath=self._keypath,
                project_id=self.project_id,
            )
        except MemstateAPIError as exc:
            if exc.status_code == 404:
                return  # Already empty
            logger.error("Failed to clear chat history in Memstate: %s", exc)
            raise


# ---------------------------------------------------------------------------
# MemstateSemanticMemory — Modern LCEL Runnable for semantic memory injection
# ---------------------------------------------------------------------------


class MemstateSemanticMemory:
    """
    Semantic memory injector for LangChain LCEL chains.

    Wraps any LCEL chain to add two memory behaviours:

    1. **Before each call** — retrieves relevant memories from Memstate via
       semantic search and injects them into the chain's input dict under
       ``memory_key`` (default: ``"relevant_memories"``).

    2. **After each call** — persists the conversation turn to Memstate via
       the server-side intelligence endpoint (``POST /api/v1/ingest``), which
       automatically extracts keypaths and creates structured memories.

    This is the modern LangChain 0.3+ replacement for the deprecated
    ``BaseMemory`` / ``ConversationChain`` pattern.

    Example::

        from memstate_langchain import MemstateSemanticMemory
        from langchain_openai import ChatOpenAI
        from langchain_core.prompts import ChatPromptTemplate

        llm = ChatOpenAI(model="gpt-4o")
        prompt = ChatPromptTemplate.from_messages([
            ("system", "You are a helpful assistant.\\n\\n{relevant_memories}"),
            ("human", "{input}"),
        ])

        chain = prompt | llm

        memory = MemstateSemanticMemory(
            api_key="mst_...",
            project_id="my-assistant",
        )

        # Wrap the chain with memory
        chain_with_memory = memory.wrap(chain)

        # Use like a normal chain
        response = chain_with_memory.invoke({"input": "My name is Alice."})
        response2 = chain_with_memory.invoke({"input": "What's my name?"})
    """

    def __init__(
        self,
        api_key: str,
        project_id: str,
        *,
        session_id: str = "default",
        base_url: str = DEFAULT_BASE_URL,
        timeout: int = DEFAULT_TIMEOUT,
        search_limit: int = DEFAULT_SEARCH_LIMIT,
        memory_key: str = "relevant_memories",
        input_key: str = "input",
        output_key: str = "content",
        human_prefix: str = "Human",
        ai_prefix: str = "AI",
        ingest_source: str = "langchain",
        wait_for_ingest: bool = False,
        ingest_poll_interval: float = DEFAULT_INGEST_POLL_INTERVAL,
        ingest_max_wait: float = DEFAULT_INGEST_MAX_WAIT,
    ) -> None:
        """
        Args:
            api_key: Memstate API key.
            project_id: Memstate project ID for namespacing memories.
            session_id: Optional session scope (used in ingest context label).
            base_url: Memstate API base URL.
            timeout: HTTP timeout in seconds.
            search_limit: Max memories to retrieve per turn.
            memory_key: Key injected into the chain's input dict with retrieved memories.
            input_key: Key in the chain's input dict containing the human message.
            output_key: Key in the chain's output to extract the AI response.
                For ChatOpenAI, this is ``"content"`` (AIMessage attribute).
            human_prefix: Label for human turns in saved summaries.
            ai_prefix: Label for AI turns in saved summaries.
            ingest_source: Source label shown in the Memstate dashboard.
            wait_for_ingest: If True, poll until async ingest completes before returning.
            ingest_poll_interval: Seconds between ingest job status polls.
            ingest_max_wait: Max seconds to wait for async ingest.
        """
        self.client = MemstateClient(api_key=api_key, base_url=base_url, timeout=timeout)
        self.project_id = project_id
        self.session_id = session_id
        self.search_limit = search_limit
        self.memory_key = memory_key
        self.input_key = input_key
        self.output_key = output_key
        self.human_prefix = human_prefix
        self.ai_prefix = ai_prefix
        self.ingest_source = ingest_source
        self.wait_for_ingest = wait_for_ingest
        self.ingest_poll_interval = ingest_poll_interval
        self.ingest_max_wait = ingest_max_wait

    def retrieve(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
        """
        Retrieve relevant memories and inject them into the input dict.

        Called automatically by the wrapped chain before each invocation.
        Returns the input dict with ``memory_key`` added.
        """
        query = inputs.get(self.input_key, "")
        if not query:
            return {**inputs, self.memory_key: ""}

        try:
            resp = self.client.search(
                query=str(query),
                project_id=self.project_id,
                limit=self.search_limit,
                keypath_prefix=MemstateClient.build_keypath(self.project_id),
            )
            results = resp.get("results", [])
            if not results:
                return {**inputs, self.memory_key: ""}

            lines = ["Relevant memories from previous sessions:"]
            for r in results:
                keypath = r.get("keypath", "")
                summary = r.get("summary", r.get("content", ""))
                score = r.get("score", 0.0)
                lines.append(f"  [{keypath}] (relevance: {score:.2f}) {summary}")

            return {**inputs, self.memory_key: "\n".join(lines)}

        except MemstateError as exc:
            logger.warning("Failed to load memories from Memstate: %s", exc)
            return {**inputs, self.memory_key: ""}

    def save(self, inputs: Dict[str, Any], output: Any) -> None:
        """
        Persist the conversation turn to Memstate.

        Called automatically by the wrapped chain after each invocation.
        Extracts the AI response from ``output`` and ingests the full turn.

        Args:
            inputs: The chain's input dict (must contain ``input_key``).
            output: The chain's output (AIMessage, dict, or string).
        """
        human_text = inputs.get(self.input_key, "")

        # Extract AI response from various output types
        if hasattr(output, "content"):
            ai_text = output.content  # AIMessage
        elif isinstance(output, dict):
            ai_text = output.get(self.output_key, str(output))
        else:
            ai_text = str(output)

        content = (
            f"## Conversation Turn\n\n"
            f"**{self.human_prefix}:** {human_text}\n\n"
            f"**{self.ai_prefix}:** {ai_text}\n"
        )

        try:
            resp = self.client.ingest(
                project_id=self.project_id,
                content=content,
                source=self.ingest_source,
                context=f"session:{self.session_id}",
            )

            if self.wait_for_ingest:
                self._wait_for_job(resp)

        except MemstateError as exc:
            logger.error("Failed to save context to Memstate: %s", exc)

    def _wait_for_job(self, ingest_resp: Dict[str, Any]) -> None:
        """Poll the job status endpoint until the ingest job completes."""
        job_id = ingest_resp.get("job_id") or ingest_resp.get("ingestion_id")
        if not job_id:
            return  # Synchronous ingest — already complete

        deadline = time.monotonic() + self.ingest_max_wait
        while time.monotonic() < deadline:
            try:
                status_resp = self.client.get_job_status(job_id)
                status = status_resp.get("status", "")
                if status in ("complete", "completed", "done"):
                    logger.debug("Ingest job %s completed", job_id)
                    return
                if status in ("error", "failed"):
                    logger.warning(
                        "Ingest job %s failed: %s",
                        job_id,
                        status_resp.get("error", "unknown error"),
                    )
                    return
            except MemstateError as exc:
                logger.debug("Error polling job %s: %s", job_id, exc)
            time.sleep(self.ingest_poll_interval)

        logger.warning(
            "Ingest job %s did not complete within %ss",
            job_id,
            self.ingest_max_wait,
        )

    def wrap(self, chain: Runnable) -> Runnable:
        """
        Wrap an LCEL chain with semantic memory retrieval and persistence.

        The returned runnable:
        1. Retrieves relevant memories and injects them into the input dict
        2. Invokes the original chain
        3. Saves the conversation turn to Memstate asynchronously

        Args:
            chain: Any LCEL ``Runnable`` (e.g., ``prompt | llm``).

        Returns:
            A new ``Runnable`` that behaves identically to ``chain`` but
            with Memstate memory retrieval and persistence.

        Example::

            memory = MemstateSemanticMemory(api_key="mst_...", project_id="app")
            chain_with_memory = memory.wrap(prompt | llm)
            response = chain_with_memory.invoke({"input": "Hello!"})
        """
        memory_self = self

        def _invoke_with_memory(inputs: Dict[str, Any]) -> Any:
            # Step 1: Retrieve relevant memories
            enriched_inputs = memory_self.retrieve(inputs)
            # Step 2: Run the chain
            output = chain.invoke(enriched_inputs)
            # Step 3: Save the conversation turn
            memory_self.save(inputs, output)
            return output

        return RunnableLambda(_invoke_with_memory)

    def get_memories(
        self,
        query: Optional[str] = None,
        *,
        limit: int = 20,
        keypath_prefix: Optional[str] = None,
    ) -> List[Dict[str, Any]]:
        """
        Directly retrieve memories from Memstate.

        If ``query`` is provided, performs semantic search.
        Otherwise, browses all memories under the project keypath.
        """
        if query:
            resp = self.client.search(
                query=query,
                project_id=self.project_id,
                limit=limit,
                keypath_prefix=keypath_prefix
                or MemstateClient.build_keypath(self.project_id),
            )
            return resp.get("results", [])
        else:
            prefix = keypath_prefix or MemstateClient.build_keypath(self.project_id)
            resp = self.client.browse(
                keypath_prefix=prefix,
                project_id=self.project_id,
                limit=limit,
            )
            return resp.get("memories", [])

    def clear(self) -> None:
        """
        Delete all memories for this project from Memstate.

        Uses a recursive soft-delete (tombstone) on the project keypath,
        preserving full version history.
        """
        keypath = MemstateClient.build_keypath(self.project_id)
        try:
            self.client.delete_memory(
                keypath=keypath,
                project_id=self.project_id,
                recursive=True,
            )
        except MemstateAPIError as exc:
            if exc.status_code == 404:
                return
            raise


# ---------------------------------------------------------------------------
# Backwards-compatible alias
# ---------------------------------------------------------------------------

# MemstateMemory is an alias for MemstateSemanticMemory for discoverability.
# Users familiar with Mem0's MemoryClient or LangChain's old memory classes
# will find this name intuitive.
MemstateMemory = MemstateSemanticMemory
