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

from typing import Any, Dict, Union

from urllib.parse import quote
from typing import TYPE_CHECKING

from pydantic import BaseModel

if TYPE_CHECKING:
    from ..client import AielClient

from ..models.memory import (
    CheckpointOut,
    CheckpointPutIn,
    RecallForgetIn,
    RecallGetIn,
    RecallPutIn,
    RecallSearchIn,
    ThreadAppendIn,
    ThreadGetOut,
    ThreadStateOut,
    ThreadStatePatchIn,
)


def _qp(segment: str) -> str:
    """
    Quote a single path segment (thread_key) safely.
    Allows ':' etc but escapes '/', '?' and spaces.
    """
    return quote(str(segment), safe=":@-._~")


class MemoryService:
    def __init__(self, client: AielClient) -> None:
        self._c = client

    # ---- Thread messages ----

    def thread_append_messages(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_key: str,
        payload: ThreadAppendIn | Dict[str, Any],
        request_id: str | None = None,
    ) -> ThreadGetOut:
        tk = _qp(thread_key)
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{tk}/messages:append"
        body = payload.model_dump(exclude_none=True) if isinstance(payload, BaseModel) else payload
        data = self._c._request("memory", "POST", path, json_body=body, request_id=request_id)
        return ThreadGetOut.model_validate(data)

    def thread_get_messages(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_key: str,
        limit: int = 50,
        request_id: str | None = None,
    ) -> ThreadGetOut:
        tk = _qp(thread_key)
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{tk}/messages"
        data = self._c._request(
            "memory",
            "GET",
            path,
            params={"limit": limit},
            request_id=request_id,
        )
        return ThreadGetOut.model_validate(data)

    # ---- Thread state ----

    def thread_get_state(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_key: str,
        request_id: str | None = None,
    ) -> ThreadStateOut:
        tk = _qp(thread_key)
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{tk}/state"
        data = self._c._request("memory", "GET", path, request_id=request_id)
        return ThreadStateOut.model_validate(data)

    def thread_patch_state(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_key: str,
        payload: ThreadStatePatchIn | Dict[str, Any],
        request_id: str | None = None,
    ) -> ThreadStateOut:
        tk = _qp(thread_key)
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{tk}/state"
        body = payload.model_dump(exclude_none=True) if isinstance(payload, BaseModel) else payload
        data = self._c._request("memory", "PUT", path, json_body=body, request_id=request_id)
        return ThreadStateOut.model_validate(data)

    # ---- Checkpoints ----

    def checkpoint_put(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_key: str,
        payload: CheckpointPutIn | Dict[str, Any],
        request_id: str | None = None,
    ) -> CheckpointOut:
        tk = _qp(thread_key)
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{tk}/checkpoint"
        body = payload.model_dump(exclude_none=True) if isinstance(payload, BaseModel) else payload
        data = self._c._request("memory", "POST", path, json_body=body, request_id=request_id)
        return CheckpointOut.model_validate(data)

    def checkpoint_get_latest(
        self,
        workspace_id: str,
        project_id: str,
        *,
        thread_key: str,
        request_id: str | None = None,
    ) -> CheckpointOut:
        tk = _qp(thread_key)
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/thread/{tk}/checkpoint/latest"
        data = self._c._request("memory", "GET", path, request_id=request_id)
        return CheckpointOut.model_validate(data)

    # ---- Recall ----

    def recall_put(
        self,
        workspace_id: str,
        project_id: str,
        *,
        payload: RecallPutIn | Dict[str, Any],
        request_id: str | None = None,
    ) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/recall:put"
        body = payload.model_dump(exclude_none=True) if isinstance(payload, BaseModel) else payload
        return self._c._request("memory", "POST", path, json_body=body, request_id=request_id) or {}

    def recall_get(
        self,
        workspace_id: str,
        project_id: str,
        *,
        payload: RecallGetIn | Dict[str, Any],
        request_id: str | None = None,
    ) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/recall:get"
        body = payload.model_dump(exclude_none=True) if isinstance(payload, BaseModel) else payload
        return self._c._request("memory", "POST", path, json_body=body, request_id=request_id) or {}

    def recall_search(
        self,
        workspace_id: str,
        project_id: str,
        *,
        payload: RecallSearchIn | Dict[str, Any],
        request_id: str | None = None,
    ) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/recall:search"
        body = payload.model_dump(exclude_none=True) if isinstance(payload, BaseModel) else payload
        return self._c._request("memory", "POST", path, json_body=body, request_id=request_id) or {}

    def recall_forget(
        self,
        workspace_id: str,
        project_id: str,
        *,
        payload: RecallForgetIn | Dict[str, Any],
        request_id: str | None = None,
    ) -> Dict[str, Any]:
        path = f"/v1/workspaces/{workspace_id}/projects/{project_id}/memory/recall:forget"
        body = payload.model_dump(exclude_none=True) if isinstance(payload, BaseModel) else payload
        return self._c._request("memory", "POST", path, json_body=body, request_id=request_id) or {}
