"""
记忆存储抽象层

hiwo-insight 不绑定任何特定的记忆存储。
你可以用：
  - InMemoryStore：内存存储（测试/演示用）
  - 自定义 MemoryStore：对接 Mem0、Memobase、PostgreSQL 等
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from datetime import datetime, timedelta
from typing import Optional

from cogprofile.models import Memory, MemoryCategory


class MemoryStore(ABC):
    """记忆存储抽象接口 - 实现这个接口即可对接任何记忆后端"""

    @abstractmethod
    async def get_memories(
        self,
        user_id: str,
        since: Optional[datetime] = None,
        until: Optional[datetime] = None,
        category: Optional[MemoryCategory] = None,
        limit: int = 500,
    ) -> list[Memory]:
        """获取指定用户的记忆列表"""
        ...

    @abstractmethod
    async def get_all_users(self) -> list[str]:
        """获取所有用户ID"""
        ...


class InMemoryStore(MemoryStore):
    """
    内存存储 - 用于测试和快速开始。

    用法:
        store = InMemoryStore()
        store.add(Memory(id="1", user_id="u1", content="我今天很开心", ...))
        memories = await store.get_memories("u1", since=datetime(2025, 1, 1))
    """

    def __init__(self):
        self._memories: list[Memory] = []

    def add(self, memory: Memory):
        """添加一条记忆"""
        self._memories.append(memory)

    def add_batch(self, memories: list[Memory]):
        """批量添加记忆"""
        self._memories.extend(memories)

    async def get_memories(
        self,
        user_id: str,
        since: Optional[datetime] = None,
        until: Optional[datetime] = None,
        category: Optional[MemoryCategory] = None,
        limit: int = 500,
    ) -> list[Memory]:
        results = [m for m in self._memories if m.user_id == user_id]

        if since:
            results = [m for m in results if m.created_at >= since]
        if until:
            results = [m for m in results if m.created_at <= until]
        if category:
            results = [m for m in results if m.category == category]

        results.sort(key=lambda m: m.created_at)
        return results[:limit]

    async def get_all_users(self) -> list[str]:
        return list(set(m.user_id for m in self._memories))


class Mem0Adapter(MemoryStore):
    """
    Mem0 适配器 - 把 Mem0 的记忆转换为 hiwo-insight 的格式。

    用法:
        from mem0 import Memory as Mem0Memory
        mem0 = Mem0Memory()
        store = Mem0Adapter(mem0_client=mem0)
        engine = InsightEngine(store=store, ...)
    """

    def __init__(self, mem0_client):
        self._client = mem0_client

    async def get_memories(
        self,
        user_id: str,
        since: Optional[datetime] = None,
        until: Optional[datetime] = None,
        category: Optional[MemoryCategory] = None,
        limit: int = 500,
    ) -> list[Memory]:
        # 从 Mem0 获取原始记忆
        raw = self._client.get_all(user_id=user_id, limit=limit)
        memories = []
        for item in raw.get("results", []):
            memories.append(Memory(
                id=item.get("id", ""),
                user_id=user_id,
                content=item.get("memory", ""),
                category=MemoryCategory.OTHER,
                created_at=datetime.fromisoformat(
                    item.get("created_at", datetime.now().isoformat())
                ),
                metadata=item.get("metadata", {}),
            ))

        if since:
            memories = [m for m in memories if m.created_at >= since]
        if until:
            memories = [m for m in memories if m.created_at <= until]

        memories.sort(key=lambda m: m.created_at)
        return memories[:limit]

    async def get_all_users(self) -> list[str]:
        # Mem0 目前没有 list_users 接口，需要自行维护
        raise NotImplementedError("Mem0 不支持列出所有用户，请自行维护用户列表")
