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

from pydantic import BaseModel


class LongTermSearchRequest(BaseModel):
    table_name: str
    query: str
    user_id: Optional[str] = None
    agent_id: Optional[str] = None
    run_id: Optional[str] = None
    limit: Optional[int] = 100
    filters: Optional[str] = None


class Data(BaseModel):
    role: str
    content: str


class LongTermAddRequest(BaseModel):
    table_name: str
    messages: Union[str, List[Data]]
    user_id: Optional[str] = None
    agent_id: Optional[str] = None
    run_id: Optional[str] = None
    metadata: Optional[str] = None
    filters: Optional[str] = None


class LongTermRemoveRequest(BaseModel):
    memory_id: str


class LongTermRemoveAllRequest(BaseModel):
    user_id: Optional[str] = None
    agent_id: Optional[str] = None
    run_id: Optional[str] = None


class LongTermUpdateRequest(BaseModel):
    memory_id: str
    data: Dict


class Content(BaseModel):
    text: str
    action: Optional[str] = None


class ShortTermAddRequest(BaseModel):
    content: Content
    user_id: str
    agent_id: Optional[str] = None
    run_id: str
    table_name: str


class ShortTermGetRequest(BaseModel):
    run_id: str
    timestamp: Optional[str] = None
    table_name: str


class ShortTermGetRunIdsRequest(BaseModel):
    timestamp: Optional[str] = None


class ShortTermDeleteRequest(BaseModel):
    id: str


class BootstrapRequest(BaseModel):
    data_type: str
    src: str
    table_name: str
    chunk_size: int = 2000
    chunk_overlap: int = 50
    min_chunk_size: int = 50
    user_id: Optional[str] = None
    agent_id: Optional[str] = None
    run_id: Optional[str] = None


class Response(BaseModel):
    message: str
    status: bool
    data: Optional[Any] = None
