"""
aiel_sdk.memory.types
====================

Stable datatypes for the AIEL Memory SDK.

This module defines small, JSON-friendly dataclasses used across saver
implementations (e.g., InMemorySaver and RemoteSaver). These types form a stable
surface for:

- thread message history (short-term memory)
- mutable thread state (workflow/runtime variables)
- checkpoints (resumable snapshots)
- recall (long-term key/value memory)
- capability metadata (feature gating / compatibility mapping)

Type aliases
------------
Json:
    A JSON-like mapping (``dict[str, Any]``). Values should be serializable.

Role:
    Allowed roles for thread messages: ``"human"``, ``"ai"``, ``"system"``, ``"tool"``.

Core concepts
-------------
thread_key:
    Stable external identifier provided by the user of the SDK. It is the primary
    handle used to read/write messages, state, and checkpoints.

thread_id:
    Backend-generated internal identifier returned by remote backends. The user does
    not provide it; it may be absent for purely in-memory implementations.

namespace:
    Hierarchical grouping for recall items, always represented as ``list[str]``
    (example: ``["tenant", "t1", "user", "u1"]``).

Notes on validation
-------------------
Some dataclasses (e.g., ThreadMessage, RecallQuery) implement light client-side
validation to fail fast before calling a remote API. Backends may still enforce
their own validation rules.
"""

from __future__ import annotations

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

Json = Dict[str, Any]
Role = Literal["human", "ai", "system", "tool"]


# ----------------------------
# Core datatypes (agnostic)
# ----------------------------

@dataclass(frozen=True)
class ThreadScope:
    """
    Optional visibility/ownership metadata for thread writes and state patches.

    ThreadScope is intended to help backends and policy layers attribute data to
    specific actors (e.g., user or agent). When omitted, backends may apply a default
    scope.

    Attributes:
        user_id:
            Optional user identifier associated with the operation.
        agent_id:
            Optional agent identifier associated with the operation.

    Methods:
        to_json():
            Return a compact JSON dict containing only fields that are not None.
    """
    user_id: Optional[str] = None
    agent_id: Optional[str] = None

    def to_json(self) -> Json:
        out: Json = {}
        if self.user_id is not None:
            out["user_id"] = self.user_id
        if self.agent_id is not None:
            out["agent_id"] = self.agent_id
        return out


@dataclass(frozen=True)
class ThreadMessage:
    """
    A single message within a thread.

    Messages are the building blocks of short-term memory (conversation history).
    They can represent user input, assistant output, system instructions, or tool
    traces.

    Attributes:
        role:
            One of: "human", "ai", "system", "tool".
        content:
            Message text. Must be non-empty after trimming and <= 20,000 characters.
        ts:
            Optional timestamp string (backend-defined format). Often set by the backend.
        meta:
            Optional JSON metadata (e.g., tool name, model info, token counts).

    Methods:
        validate():
            Raises ValueError for malformed content (empty or too large).
        to_json():
            Validate and return a JSON representation suitable for API payloads.
    """
    role: Role
    content: str
    ts: Optional[str] = None
    meta: Json = field(default_factory=dict)

    def validate(self) -> None:
        # Fail-fast client-side; backend will also validate.
        c = (self.content or "").strip()
        if not c:
            raise ValueError("ThreadMessage.content must be non-empty")
        if len(c) > 20000:
            raise ValueError("ThreadMessage.content must be <= 20000 characters")

    def to_json(self) -> Json:
        self.validate()
        out: Json = {"role": self.role, "content": self.content, "meta": dict(self.meta or {})}
        if self.ts is not None:
            out["ts"] = self.ts
        return out


@dataclass(frozen=True)
class Thread:
    """
    Thread message container returned by saver read/write operations.

    Attributes:
        thread_key:
            Stable external identifier for the conversation/session.
        thread_id:
            Backend-generated internal identifier (may be None for in-memory backends).
        messages:
            Ordered list of thread messages (often most-recent-first depending on backend).
        updated_at:
            Optional last-updated timestamp (backend-defined format).
    """
    # External stable handle users pass everywhere
    thread_key: str
    # Backend-derived deterministic id (may be absent for in-memory)
    thread_id: Optional[str] = None
    messages: List[ThreadMessage] = field(default_factory=list)
    updated_at: Optional[str] = None

@dataclass(frozen=True)
class ThreadStateIn:
    """
    Input payload shape for patching thread state.

    This type is useful when you want a structured container for state patch
    operations (for example, passing through SDK layers).

    Attributes:
        thread_key:
            Stable external identifier for the conversation/session.
        patch:
            JSON-like dict representing state updates (backend defines merge semantics).
        scope:
            Optional ThreadScope for attribution/visibility.
        ttl_seconds:
            Optional time-to-live for the patched state, if supported by backend.
    """
    thread_key: str
    patch: Json = field(default_factory=dict)
    scope: ThreadScope | None = None,
    ttl_seconds: Optional[int] | None = None,

@dataclass(frozen=True)
class ThreadState:
    """
    Structured state associated with a thread.

    State is intended for mutable workflow/runtime variables (e.g., routing decisions,
    extracted entities, progress markers). It is separate from message history.

    Attributes:
        thread_key:
            Stable external identifier for the conversation/session.
        thread_id:
            Backend-generated internal identifier (may be None for in-memory backends).
        state:
            JSON-like dict representing the current thread state.
        updated_at:
            Optional last-updated timestamp (backend-defined format).
    """

    thread_key: str
    thread_id: Optional[str] = None
    state: Json = field(default_factory=dict)
    updated_at: Optional[str] = None

@dataclass(frozen=True)
class Checkpoint:
    """
    A point-in-time snapshot of thread state for recovery/resume.

    Checkpoints are typically written at key milestones so workflows/agents can resume
    without replaying the full conversation or rebuilding state from scratch.

    Attributes:
        thread_key:
            Stable external identifier for the conversation/session.
        thread_id:
            Backend-generated internal identifier (may be None for in-memory backends).
        checkpoint_id:
            Backend-generated checkpoint identifier, when available.
        state:
            JSON-like dict representing the saved state snapshot.
        metadata:
            JSON-like dict describing the checkpoint (e.g., node name, source agent, version).
        ttl_seconds:
            Optional time-to-live for the checkpoint, if supported by backend.
        created_at:
            Optional creation timestamp (backend-defined format).
    """
    thread_key: str
    thread_id: Optional[str] = None
    checkpoint_id: Optional[str] = None
    state: Json = field(default_factory=dict)
    metadata: Json = field(default_factory=dict)
    ttl_seconds: Optional[int] | None = None
    created_at: Optional[str] = None


@dataclass(frozen=True)
class RecallItem:
    """
    A long-term memory item stored under a hierarchical namespace.

    Recall is a durable key/value store intended for facts and preferences
    (e.g., user settings, profile attributes, stable business context).

    Attributes:
        namespace:
            Hierarchical namespace as list[str] (example: ["users", "42"]).
        key:
            Key within the namespace (example: "preferences", "profile:v2").
        value:
            JSON-like dict payload.
        tags:
            Optional JSON tags for filtering/analytics (backend-dependent).
        updated_at:
            Optional last-updated timestamp (backend-defined format).
    """
    namespace: List[str]
    key: str
    value: Json = field(default_factory=dict)
    tags: Json = field(default_factory=dict)
    updated_at: Optional[str] = None

@dataclass(frozen=True)
class RecallResult:
    """
    Recall lookup result wrapper.

    Some backends return a wrapper indicating whether the item was found, plus the
    item payload. This type normalizes that shape.

    Attributes:
        found:
            True if the recall entry exists.
        item:
            JSON-like dict containing the stored item payload (backend-defined shape).
    """
    found: bool
    item: Json = field(default_factory=dict)

@dataclass(frozen=True)
class RecallPutResult:
    """
    Recall put result wrapper.

    Some backends return a wrapper indicating whether the write succeeded, plus an
    optional document id. This type normalizes that shape.

    Attributes:
        ok:
            True if the put operation succeeded.
        id:
            Optional document id returned by the backend.
    """
    ok: bool
    id: str | None = field(default=None)

@dataclass(frozen=True)
class RecallQuery:
    """
    Search request for recall items within a namespace.

    Attributes:
        namespace:
            Hierarchical namespace as list[str].
        query:
            JSON-like filter object (backend-defined semantics).
        limit:
            Maximum number of results to return. Must be between 1 and 200.

    Methods:
        validate():
            Raises ValueError when limit is out of bounds.
    """
    namespace: List[str]
    query: Json = field(default_factory=dict)
    limit: int = 20

    def validate(self) -> None:
        if self.limit < 1 or self.limit > 200:
            raise ValueError("RecallQuery.limit must be between 1 and 200")


# ----------------------------
# Capability metadata (stable surface)
# ----------------------------

@dataclass(frozen=True)
class SaverCapability:
    """
    A single capability exposed by a saver implementation.

    Capabilities are used for feature gating and compatibility mapping, especially
    when backends enforce "allowed_actions" style permissions.

    Attributes:
        id:
            Stable capability identifier (example: "memory:thread:read").
        label:
            Human-friendly label.
        description:
            Short description of the capability.
        group:
            Group key used for organizing related capabilities.
        group_label:
            Human-friendly group label.
    """

    id: str
    label: str
    description: str
    group: str
    group_label: str


@dataclass(frozen=True)
class SaverCapabilities:
    """
    Collection of saver capabilities.

    This is intended as a stable product surface (not UI fluff) used for:
    - feature gating
    - backend 'allowed_actions' compatibility mapping

    Attributes:
        items:
            Tuple of SaverCapability entries.

    Methods:
        ids():
            Return a list of capability IDs in the same order as items.
    """
    items: Tuple[SaverCapability, ...]

    def ids(self) -> List[str]:
        return [c.id for c in self.items]


DEFAULT_CAPABILITIES = SaverCapabilities(
    items=[
        # ---- Thread (short-term) ----
        SaverCapability(
            id="memory:thread:read",
            label="Read Thread Memory",
            description="Read conversation messages from thread memory (short-term).",
            group="Short-term",
            group_label="Project memory integration",
        ),
        SaverCapability(
            id="memory:thread:write",
            label="Write Thread Memory",
            description="Append messages to thread memory (short-term).",
            group="Short-term",
            group_label="Project memory integration",
        ),
        # ---- State (working state) ----
        SaverCapability(
            id="memory:state:read",
            label="Read Thread State",
            description="Read the structured state associated with a thread (workflow variables).",
            group="Short-term",
            group_label="Project memory integration",
        ),
        SaverCapability(
            id="memory:state:write",
            label="Write Thread State",
            description="Patch/update the structured state associated with a thread.",
            group="Short-term",
            group_label="Project memory integration",
        ),
        # ---- Checkpoints (resume points) ----
        SaverCapability(
            id="memory:checkpoint:read",
            label="Read Checkpoints",
            description="Read latest checkpoint for a thread (resume/recovery).",
            group="Checkpoints",
            group_label="Project memory integration",
        ),
        SaverCapability(
            id="memory:checkpoint:write",
            label="Write Checkpoints",
            description="Create checkpoints for a thread (resume/retry support).",
            group="Checkpoints",
            group_label="Project memory integration",
        ),
        # ---- Recall (long-term) ----
        SaverCapability(
            id="memory:recall:read",
            label="Read Long-term Memory",
            description="Read recall items (get/search) from long-term memory.",
            group="Recall",
            group_label="Project memory integration",
        ),
        SaverCapability(
            id="memory:recall:write",
            label="Write Long-term Memory",
            description="Write recall items (put) to long-term memory.",
            group="Recall",
            group_label="Project memory integration",
        ),
        SaverCapability(
            id="memory:recall:delete",
            label="Delete Long-term Memory",
            description="Delete recall items (forget) from long-term memory.",
            group="Recall",
            group_label="Project memory integration",
        ),
    ]
)