# 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
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):
    """
    Local saver aligned with backend semantics.
    Backend-derived IDs don't exist here; thread_id/checkpoint_id remain None.
    """

    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:
        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:
        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:
        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:
        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]:
        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:
        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,
    ) -> None:
        self._recall[(tuple(namespace), key)] = RecallItem(
            namespace=list(namespace), key=key, value=dict(value), tags=dict(tags or {})
        )

    def recall_get(self, *, namespace: List[str], key: str) -> Optional[RecallItem]:
        return self._recall.get((tuple(namespace), key))

    def recall_search(self, *, query: RecallQuery) -> List[RecallItem]:
        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:
        key = selector.get("key")
        if isinstance(key, str) and key:
            self._recall.pop((tuple(namespace), key), None)