# aiel_sdk/memory/in_memory.py
from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Sequence, Tuple

from .saver import Saver
from .types import Checkpoint, RecallItem, RecallQuery, Thread, ThreadMessage, ThreadScope, ThreadState, RecallPutResult
from .thread_key import normalize_thread_key  # if you kept it; else inline normalize


@dataclass
class _InMemThread:
    messages: List[ThreadMessage] = field(default_factory=list)
    state: Dict[str, Any] = field(default_factory=dict)
    checkpoints: List[Checkpoint] = field(default_factory=list)


class InMemorySaver(Saver):
    """
        InMemorySaver
        =============

        In-process, non-persistent saver for local development and tests.

        InMemorySaver mirrors the public behavior of the remote Memory API surface
        (threads, state, checkpoints, recall) but stores everything in Python data
        structures inside the current process.

        When to use
        -----------
        - Local prototyping
        - Deterministic behavior without network calls
        - CI pipelines where you want fast, side-effect-free runs

        When *not* to use
        -----------------
        - Production or any scenario requiring persistence across process restarts
        - Multi-instance deployments (state is not shared between processes)

        Behavioral notes
        ----------------
        - thread_key is normalized via :func:`~aiel_sdk.memory.thread_key.normalize_thread_key`.
        - Backend-derived IDs do not exist here:
            - Thread.thread_id is always None
            - Checkpoint.checkpoint_id is always None
        - ttl_seconds is accepted for API symmetry but is not enforced.
        (If you need TTL behavior in tests, wrap this saver or add a small eviction layer.)

        Data model
        ----------
        For each normalized thread_key this saver maintains:
        - messages: list[ThreadMessage]
        - state: dict[str, Any] (shallow-merged updates)
        - checkpoints: list[Checkpoint]

        Recall is stored separately as a mapping keyed by (namespace tuple, key).

        Thread API
        ----------
        thread_write(...):
            Appends messages to the thread in insertion order.

        thread_read(...):
            Returns the most recent `limit` messages (if limit > 0).
            If limit <= 0, returns all messages.

        State API
        ---------
        state_patch(...):
            Performs a shallow merge: ``state.update(patch)``.
            This is intentionally predictable and backend-agnostic.

        Checkpoints API
        ---------------
        checkpoint_write(...):
            Stores a snapshot of the provided state + metadata.
            Returns the stored Checkpoint.

        checkpoint_read_latest(...):
            Returns the most recently written checkpoint or None.

        Recall API
        ----------
        recall_put(...):
            Stores/overwrites a RecallItem for (namespace, key).

        recall_get(...):
            Returns the RecallItem for (namespace, key) or None.

        recall_search(...):
            Returns up to `query.limit` items whose namespace starts with the query namespace.

        recall_forget(...):
            Deletes a single item when selector contains a non-empty string under "key".

        Examples
        --------
            from aiel_sdk.memory import InMemorySaver, ThreadMessage, RecallQuery

            saver = InMemorySaver()
            thread_key = "user:42/session:abc"

            saver.thread_write(
                thread_key=thread_key,
                messages=[ThreadMessage(role="human", content="Hello!")],
            )
            thread = saver.thread_read(thread_key=thread_key, limit=50)

            saver.state_patch(thread_key=thread_key, patch={"step": "triage"})
            state = saver.state_read(thread_key=thread_key)

            cp = saver.checkpoint_write(thread_key=thread_key, state=state.state, metadata={"node": "triage"})

            saver.recall_put(namespace=["users", "42"], key="preferences", value={"tone": "concise"})
            item = saver.recall_get(namespace=["users", "42"], key="preferences")
            items = saver.recall_search(query=RecallQuery(namespace=["users", "42"], limit=20))
        """

    def __init__(self) -> None:
        self._threads: Dict[str, _InMemThread] = {}
        self._recall: Dict[Tuple[Tuple[str, ...], str], RecallItem] = {}

    def _t(self, thread_key: str) -> tuple[str, _InMemThread]:
        k = normalize_thread_key(thread_key)
        if k not in self._threads:
            self._threads[k] = _InMemThread()
        return k, self._threads[k]

    # ---- Thread ----
    def thread_read(self, *, thread_key: str, limit: int = 50) -> Thread:
        """
        Read messages for a thread (short-term memory).

        This returns a :class:`~aiel_sdk.memory.types.Thread` containing the messages
        currently stored for the given ``thread_key``.

        Behavior (InMemorySaver)
        ------------------------
        - ``thread_key`` is normalized before lookup/creation.
        - If the thread does not exist yet, it is created empty and an empty thread is returned.
        - If ``limit > 0``, returns only the most recent ``limit`` messages.
        - If ``limit <= 0``, returns all messages.
        - ``thread_id`` is always ``None`` and ``updated_at`` is always ``None`` in this saver.

        Args:
            thread_key:
                Stable external thread identifier (will be normalized).
            limit:
                Maximum number of messages to return. Defaults to 50.
                Use a bounded limit to keep payloads small.

        Returns:
            Thread:
                Thread container with ``thread_key``, ``messages``, and no backend IDs.
        """
        k, t = self._t(thread_key)
        msgs = t.messages[-limit:] if limit > 0 else list(t.messages)
        return Thread(thread_key=k, thread_id=None, messages=list(msgs), updated_at=None)

    def thread_write(
        self,
        *,
        thread_key: str,
        messages: Sequence[ThreadMessage],
        scope: ThreadScope | None = None,
        ttl_seconds: int | None = None,
    ) -> Thread:
        """
        Append messages to a thread (short-term memory).

        Behavior (InMemorySaver)
        ------------------------
        - ``thread_key`` is normalized before lookup/creation.
        - If the thread does not exist yet, it is created.
        - Messages are appended in insertion order.
        - ``scope`` and ``ttl_seconds`` are accepted for API symmetry but are not enforced.
        - ``thread_id`` and ``updated_at`` are always ``None``.

        Args:
            thread_key:
                Stable external thread identifier (will be normalized).
            messages:
                Messages to append to the thread.
            scope:
                Optional attribution/visibility metadata. Not enforced in-memory.
            ttl_seconds:
                Optional TTL. Not enforced in-memory.

        Returns:
            Thread:
                Updated thread with all stored messages.
        """
        k, t = self._t(thread_key)
        t.messages.extend(list(messages))
        return Thread(thread_key=k, thread_id=None, messages=list(t.messages), updated_at=None)

    # ---- State ----
    def state_read(self, *, thread_key: str) -> ThreadState:
        """
        Read the structured state associated with a thread.

        Behavior (InMemorySaver)
        ------------------------
        - ``thread_key`` is normalized before lookup/creation.
        - If the thread does not exist yet, it is created with empty state.
        - Returns a copy of the stored state dict.
        - ``thread_id`` and ``updated_at`` are always ``None``.

        Args:
            thread_key:
                Stable external thread identifier (will be normalized).

        Returns:
            ThreadState:
                Container holding the current state dict for the thread.
        """
        k, t = self._t(thread_key)
        return ThreadState(thread_key=k, thread_id=None, state=dict(t.state), updated_at=None)

    def state_patch(
        self,
        *,
        thread_key: str,
        patch: Dict[str, Any],
        scope: ThreadScope | None = None,
        ttl_seconds: int | None = None,
    ) -> ThreadState:
        """
        Patch (shallow-merge) the structured state for a thread.

        Behavior (InMemorySaver)
        ------------------------
        - ``thread_key`` is normalized before lookup/creation.
        - If the thread does not exist yet, it is created with empty state.
        - Applies a shallow merge: ``state.update(patch)``.
        This means nested dicts are replaced, not deep-merged.
        - ``scope`` and ``ttl_seconds`` are accepted for API symmetry but are not enforced.
        - ``thread_id`` and ``updated_at`` are always ``None``.

        Args:
            thread_key:
                Stable external thread identifier (will be normalized).
            patch:
                Dict of state fields to update.
            scope:
                Optional attribution/visibility metadata. Not enforced in-memory.
            ttl_seconds:
                Optional TTL. Not enforced in-memory.

        Returns:
            ThreadState:
                Updated state container after applying the patch.
        """
        k, t = self._t(thread_key)
        t.state.update(dict(patch or {}))  # shallow merge, predictable
        return ThreadState(thread_key=k, thread_id=None, state=dict(t.state), updated_at=None)

    # ---- Checkpoints ----
    def checkpoint_read_latest(self, *, thread_key: str) -> Optional[Checkpoint]:
        """
        Read the latest checkpoint for a thread, if any exists.

        Behavior (InMemorySaver)
        ------------------------
        - ``thread_key`` is normalized before lookup/creation.
        - If the thread does not exist yet, it is created with no checkpoints.
        - Returns the most recently written checkpoint or None.
        - ``checkpoint_id`` is always ``None`` (no backend IDs).

        Args:
            thread_key:
                Stable external thread identifier (will be normalized).

        Returns:
            Checkpoint | None:
                The latest checkpoint or None if none have been written.
        """
        k, t = self._t(thread_key)
        return t.checkpoints[-1] if t.checkpoints else None

    def checkpoint_write(
        self,
        *,
        thread_key: str,
        state: Dict[str, Any],
        metadata: Dict[str, Any] | None = None,
        scope: ThreadScope | None = None,
        ttl_seconds: int | None = None,
    ) -> Checkpoint:
        """
        Write a checkpoint snapshot for a thread.

        Checkpoints are point-in-time snapshots intended for recovery/resume.

        Behavior (InMemorySaver)
        ------------------------
        - ``thread_key`` is normalized before lookup/creation.
        - If the thread does not exist yet, it is created.
        - Stores a snapshot of ``state`` (copied) plus ``metadata`` (copied).
        - ``scope`` and ``ttl_seconds`` are accepted for API symmetry but are not enforced.
        - ``thread_id`` and ``checkpoint_id`` are always ``None``; ``created_at`` is ``None``.

        Args:
            thread_key:
                Stable external thread identifier (will be normalized).
            state:
                Full state snapshot to store.
            metadata:
                Optional metadata describing the checkpoint (e.g., node name, source).
            scope:
                Optional attribution/visibility metadata. Not enforced in-memory.
            ttl_seconds:
                Optional TTL. Not enforced in-memory.

        Returns:
            Checkpoint:
                The stored checkpoint snapshot.
        """
        k, t = self._t(thread_key)
        cp = Checkpoint(thread_key=k, thread_id=None, checkpoint_id=None, state=dict(state), metadata=dict(metadata or {}))
        t.checkpoints.append(cp)
        return cp

    # ---- Recall ----
    def recall_put(
        self,
        *,
        namespace: List[str],
        key: str,
        value: Dict[str, Any],
        tags: Dict[str, Any] | None = None,
        ttl_seconds: int | None = None,
    ) -> RecallPutResult:
        """
        Store/overwrite a long-term recall item.

        Behavior (InMemorySaver)
        ------------------------
        - Stores a :class:`~aiel_sdk.memory.types.RecallItem` keyed by (namespace, key).
        - Overwrites any existing item with the same (namespace, key).
        - ``tags`` is stored as provided (copied).
        - ``ttl_seconds`` is accepted for API symmetry but is not enforced.

        Args:
            namespace:
                Hierarchical namespace as list[str] (example: ["users", "42"]).
            key:
                Item key within the namespace.
            value:
                JSON-like dict payload to store.
            tags:
                Optional tags for filtering/analytics (not interpreted in-memory).
            ttl_seconds:
                Optional TTL. Not enforced in-memory.

        """
        self._recall[(tuple(namespace), key)] = RecallItem(
            namespace=list(namespace),
            key=key,
            value=dict(value or {}),
            tags=dict(tags or {}),
        )
        return RecallPutResult(ok=True,id="fake|doc_id")

    def recall_get(self, *, namespace: List[str], key: str) -> Optional[RecallItem]:
        """
        Read a recall item by namespace + key.

        Args:
            namespace:
                Hierarchical namespace as list[str].
            key:
                Recall key within the namespace.

        Returns:
            RecallItem | None:
                The stored item, or None if not found.
        """
        return self._recall.get((tuple(namespace), key))

    def recall_search(self, *, query: RecallQuery) -> List[RecallItem]:
        """
        Search recall items within a namespace prefix.

        Behavior (InMemorySaver)
        ------------------------
        - Returns items whose namespace starts with ``query.namespace`` (prefix match).
        - Stops after returning ``query.limit`` items.
        - ``query.query`` is ignored (no filtering beyond namespace prefix).
        - Order is insertion order of the underlying mapping (implementation detail).
        Do not rely on ordering for correctness.

        Args:
            query:
                RecallQuery containing:
                - namespace: prefix namespace to match
                - query: ignored by InMemorySaver
                - limit: maximum number of items to return

        Returns:
            list[RecallItem]:
                Matching recall items (possibly empty).
        """
        ns = tuple(query.namespace)
        out: List[RecallItem] = []
        for (k_ns, _k), item in self._recall.items():
            if k_ns[: len(ns)] == ns:
                out.append(item)
                if len(out) >= query.limit:
                    break
        return out

    def recall_forget(self, *, namespace: List[str], selector: Dict[str, Any]) -> None:
        """
        Delete recall items within a namespace using a simple selector.

        Behavior (InMemorySaver)
        ------------------------
        - Supports deleting a single item when selector contains:
            {"key": "<non-empty str>"}
        - Other selector shapes are ignored (no-op).

        Args:
            namespace:
                Hierarchical namespace as list[str].
            selector:
                Deletion selector. Currently supports only "key".

        Returns:
            None
        """
        key = selector.get("key")
        if isinstance(key, str) and key:
            self._recall.pop((tuple(namespace), key), None)
        
        return {
            "deleted": len(key)
        }
