from datetime import datetime
from typing import Protocol, TypeVar
from uuid import UUID

from memv.models import Episode, Message, SemanticKnowledge

T = TypeVar("T")


class MessageStore(Protocol):
    async def add(self, message: Message) -> None: ...
    async def get(self, message_id: UUID | str) -> Message | None: ...
    async def get_by_user(self, user_id: str) -> list[Message]: ...
    async def get_by_time_range(self, user_id: str, start: datetime, end: datetime) -> list[Message]: ...


class EpisodeStore(Protocol):
    async def add(self, episode: Episode) -> None: ...
    async def get(self, episode_id: UUID | str) -> Episode | None: ...
    async def get_by_user(self, user_id: str) -> list[Episode]: ...
    async def get_by_time_range(self, user_id: str, start_time: datetime, end_time: datetime) -> list[Episode]: ...


class KnowledgeStore(Protocol):
    async def add(self, knowledge: SemanticKnowledge) -> None: ...
    async def get(self, knowledge_id: UUID | str) -> SemanticKnowledge | None: ...
    async def get_by_episode(self, episode_id: UUID | str) -> list[SemanticKnowledge]: ...


class EmbeddingClient(Protocol):
    async def embed(self, text: str) -> list[float]: ...
    async def embed_batch(self, texts: list[str]) -> list[list[float]]: ...


class LLMClient(Protocol):
    """Protocol for LLM completions with structured output."""

    async def generate(self, prompt: str) -> str: ...
    async def generate_structured(self, prompt: str, response_model: type[T]) -> T: ...
