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

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

from pydantic import BaseModel, ConfigDict, Field


# ----------------------------
# Shared
# ----------------------------

Role = Literal["human", "ai", "system", "tool"]


class ThreadScope(BaseModel):
    model_config = ConfigDict(extra="allow")
    user_id: Optional[str] = None
    agent_id: Optional[str] = None


class ChatMessageIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    role: Role
    content: str = Field(..., min_length=1, max_length=20000)
    ts: Optional[str] = None
    meta: Dict[str, Any] = Field(default_factory=dict)


# ----------------------------
# Thread messages
# ----------------------------

class ThreadAppendIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    scope: ThreadScope = Field(default_factory=ThreadScope)
    messages: List[ChatMessageIn] = Field(..., min_items=1, max_items=100)
    ttl_seconds: Optional[int] = Field(default=None, ge=60, le=31536000)


class ThreadGetOut(BaseModel):
    model_config = ConfigDict(extra="allow")
    thread_id: str
    messages: List[Dict[str, Any]] = Field(default_factory=list)
    updated_at: str


# ----------------------------
# Thread state
# ----------------------------

class ThreadStatePatchIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    scope: ThreadScope = Field(default_factory=ThreadScope)
    patch: Dict[str, Any] = Field(default_factory=dict)
    ttl_seconds: Optional[int] = Field(default=None, ge=60, le=31536000)


class ThreadStateOut(BaseModel):
    model_config = ConfigDict(extra="allow")
    thread_id: str
    state: Dict[str, Any] = Field(default_factory=dict)
    updated_at: str


# ----------------------------
# Checkpoints
# ----------------------------

class CheckpointPutIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    scope: ThreadScope = Field(default_factory=ThreadScope)
    state: Dict[str, Any] = Field(default_factory=dict)
    metadata: Dict[str, Any] = Field(default_factory=dict)
    ttl_seconds: Optional[int] = Field(default=None, ge=60, le=31536000)


class CheckpointOut(BaseModel):
    model_config = ConfigDict(extra="allow")
    thread_id: str
    checkpoint_id: str
    state: Dict[str, Any] = Field(default_factory=dict)
    metadata: Dict[str, Any] = Field(default_factory=dict)
    created_at: str


# ----------------------------
# Recall (long-term)
# ----------------------------

class RecallPutIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    namespace: List[str] = Field(..., min_items=1, max_items=16)
    key: str = Field(..., min_length=1, max_length=256)
    value: Dict[str, Any] = Field(default_factory=dict)
    tags: Dict[str, Any] = Field(default_factory=dict)
    ttl_seconds: Optional[int] = Field(default=None, ge=60, le=31536000)


class RecallGetIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    namespace: List[str] = Field(..., min_items=1, max_items=16)
    key: str = Field(..., min_length=1, max_length=256)


class RecallSearchIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    namespace: List[str] = Field(..., min_items=1, max_items=16)
    query: Dict[str, Any] = Field(default_factory=dict)
    limit: int = Field(default=20, ge=1, le=200)


class RecallForgetIn(BaseModel):
    model_config = ConfigDict(extra="allow")
    namespace: List[str] = Field(..., min_items=1, max_items=16)
    selector: Dict[str, Any] = Field(default_factory=dict)