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

from typing import Any, Dict, List, Optional, Sequence

from ..client import AielClient

from .saver import Saver
from .types import Checkpoint, RecallItem, RecallQuery, Thread, ThreadMessage, ThreadScope, ThreadState
from .thread_key import normalize_thread_key


def _parse_messages(raw_messages: List[Dict[str, Any]] | None) -> List[ThreadMessage]:
    out: List[ThreadMessage] = []
    for m in (raw_messages or []):
        if not isinstance(m, dict):
            continue
        role = str(m.get("role") or "human")
        if role not in ("human", "ai", "system", "tool"):
            role = "human"
        out.append(
            ThreadMessage(
                role=role,  # type: ignore[arg-type]
                content=str(m.get("content") or ""),
                ts=m.get("ts"),
                meta=dict(m.get("meta") or {}),
            )
        )
    return out


class RemoteSaver(Saver):
    """
    Remote saver bound to workspace/project.
    Uses MemoryService(thread_key in path); backend derives thread_id.
    """

    def __init__(self, client: AielClient, *, workspace_id: str, project_id: str) -> None:
        self._c = client
        self._ws = workspace_id
        self._proj = project_id

    # ---- Thread ----
    def thread_read(self, *, thread_key: str, limit: int = 50) -> Thread:
        k = normalize_thread_key(thread_key)
        out = self._c.memory.thread_get_messages(self._ws, self._proj, thread_key=k, limit=limit)
        return Thread(thread_key=k, thread_id=out.thread_id, messages=_parse_messages(out.messages), updated_at=out.updated_at)

    def thread_write(
        self,
        *,
        thread_key: str,
        messages: Sequence[ThreadMessage],
        scope: ThreadScope | None = None,
        ttl_seconds: int | None = None,
    ) -> Thread:
        k = normalize_thread_key(thread_key)
        payload = {
            "scope": (scope or ThreadScope()).to_json(),
            "messages": [m.to_json() for m in messages],
            "ttl_seconds": ttl_seconds,
        }
        out = self._c.memory.thread_append_messages(self._ws, self._proj, thread_key=k, payload=payload)
        return Thread(thread_key=k, thread_id=out.thread_id, messages=_parse_messages(out.messages), updated_at=out.updated_at)

    # ---- State ----
    def state_read(self, *, thread_key: str) -> ThreadState:
        k = normalize_thread_key(thread_key)
        out = self._c.memory.thread_get_state(self._ws, self._proj, thread_key=k)
        return ThreadState(thread_key=k, thread_id=out.thread_id, state=dict(out.state or {}), updated_at=out.updated_at)

    def state_patch(
        self,
        *,
        thread_key: str,
        patch: Dict[str, Any],
        scope: ThreadScope | None = None,
        ttl_seconds: int | None = None,
    ) -> ThreadState:
        k = normalize_thread_key(thread_key)
        payload = {
            "scope": (scope or ThreadScope()).to_json(),
            "patch": patch or {},
            "ttl_seconds": ttl_seconds,
        }
        out = self._c.memory.thread_patch_state(self._ws, self._proj, thread_key=k, payload=payload)
        return ThreadState(thread_key=k, thread_id=out.thread_id, state=dict(out.state or {}), updated_at=out.updated_at)

    # ---- Checkpoints ----
    def checkpoint_read_latest(self, *, thread_key: str) -> Optional[Checkpoint]:
        k = normalize_thread_key(thread_key)
        raw = self._c.memory.checkpoint_get_latest(self._ws, self._proj, thread_key=k) or {}
        # If your endpoint returns 404/no-body, MemoryService should handle that.
        if not raw:
            return None
        # If you switched MemoryService to return CheckpointOut, then 'raw' is a model.
        if hasattr(raw, "model_dump"):
            d = raw.model_dump()
        else:
            d = dict(raw)
        return Checkpoint(
            thread_key=k,
            thread_id=d.get("thread_id"),
            checkpoint_id=d.get("checkpoint_id"),
            state=dict(d.get("state") or {}),
            metadata=dict(d.get("metadata") or {}),
            created_at=d.get("created_at"),
        )

    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 = normalize_thread_key(thread_key)
        payload = {
            "scope": (scope or ThreadScope()).to_json(),
            "state": state or {},
            "metadata": metadata or {},
            "ttl_seconds": ttl_seconds,
        }
        raw = self._c.memory.checkpoint_put(self._ws, self._proj, thread_key=k, payload=payload) or {}
        if hasattr(raw, "model_dump"):
            d = raw.model_dump()
        else:
            d = dict(raw)
        return Checkpoint(
            thread_key=k,
            thread_id=d.get("thread_id"),
            checkpoint_id=d.get("checkpoint_id"),
            state=dict(d.get("state") or state or {}),
            metadata=dict(d.get("metadata") or metadata or {}),
            created_at=d.get("created_at"),
        )

    # ---- 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._c.memory.recall_put(
            self._ws,
            self._proj,
            payload={"namespace": namespace, "key": key, "value": value or {}, "tags": tags or {}, "ttl_seconds": ttl_seconds},
        )

    def recall_get(self, *, namespace: List[str], key: str) -> Optional[RecallItem]:
        raw = self._c.memory.recall_get(self._ws, self._proj, payload={"namespace": namespace, "key": key}) or {}
        if not raw:
            return None
        return RecallItem(
            namespace=list(raw.get("namespace") or namespace),
            key=str(raw.get("key") or key),
            value=dict(raw.get("value") or {}),
            tags=dict(raw.get("tags") or {}),
            updated_at=raw.get("updated_at"),
        )

    def recall_search(self, *, query: RecallQuery) -> List[RecallItem]:
        raw = self._c.memory.recall_search(
            self._ws,
            self._proj,
            payload={"namespace": query.namespace, "query": query.query or {}, "limit": query.limit},
        ) or {}
        items = raw.get("items") if isinstance(raw, dict) else None
        if not isinstance(items, list):
            return []
        out: List[RecallItem] = []
        for it in items:
            if not isinstance(it, dict):
                continue
            out.append(
                RecallItem(
                    namespace=list(it.get("namespace") or query.namespace),
                    key=str(it.get("key") or ""),
                    value=dict(it.get("value") or {}),
                    tags=dict(it.get("tags") or {}),
                    updated_at=it.get("updated_at"),
                )
            )
        return out

    def recall_forget(self, *, namespace: List[str], selector: Dict[str, Any]) -> None:
        self._c.memory.recall_forget(self._ws, self._proj, payload={"namespace": namespace, "selector": selector or {}})
