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

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

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


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

@dataclass(frozen=True)
class ThreadScope:
    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:
    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:
    # 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 ThreadState:
    thread_key: str
    thread_id: Optional[str] = None
    state: Json = field(default_factory=dict)
    updated_at: Optional[str] = None


@dataclass(frozen=True)
class Checkpoint:
    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)
    created_at: Optional[str] = None


@dataclass(frozen=True)
class RecallItem:
    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 RecallQuery:
    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:
    id: str
    label: str
    description: str
    group: str
    group_label: str


@dataclass(frozen=True)
class SaverCapabilities:
    items: Sequence[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",
        ),
    ]
)