"""Working Memory - 基于 Redis 的短期任务记忆"""

import json
import logging
from datetime import datetime, timezone
from typing import Any

import redis.asyncio as aioredis

from agentkit.memory.base import Memory, MemoryItem

logger = logging.getLogger(__name__)


class WorkingMemory(Memory):
    """Working Memory - 当前任务的上下文和中间状态

    基于 Redis 实现，支持自动过期（TTL）。
    生命周期：单次任务。
    """

    def __init__(
        self,
        redis: aioredis.Redis,
        key_prefix: str = "agentkit:working",
        default_ttl: int = 3600,
    ):
        self._redis = redis
        self._key_prefix = key_prefix
        self._default_ttl = default_ttl

    def _make_key(self, key: str) -> str:
        return f"{self._key_prefix}:{key}"

    async def store(self, key: str, value: Any, metadata: dict[str, Any] | None = None) -> None:
        redis_key = self._make_key(key)
        item = MemoryItem(
            key=key,
            value=value,
            metadata=metadata or {},
            created_at=datetime.now(timezone.utc),
        )
        await self._redis.setex(
            redis_key,
            self._default_ttl,
            json.dumps(item.to_dict(), default=str),
        )

    async def retrieve(self, key: str) -> MemoryItem | None:
        redis_key = self._make_key(key)
        data = await self._redis.get(redis_key)
        if data is None:
            return None
        item_dict = json.loads(data)
        return MemoryItem(
            key=item_dict["key"],
            value=item_dict["value"],
            metadata=item_dict.get("metadata", {}),
            score=item_dict.get("score", 1.0),
            created_at=datetime.fromisoformat(item_dict["created_at"]) if item_dict.get("created_at") else datetime.now(timezone.utc),
        )

    async def search(self, query: str, top_k: int = 5, filters: dict[str, Any] | None = None) -> list[MemoryItem]:
        """Working Memory 不支持语义检索，按 key 前缀匹配"""
        pattern = self._make_key(f"{query}*")
        keys = []
        async for key in self._redis.scan_iter(match=pattern, count=top_k):
            keys.append(key)
            if len(keys) >= top_k:
                break

        items = []
        for key in keys:
            data = await self._redis.get(key)
            if data:
                item_dict = json.loads(data)
                items.append(MemoryItem(
                    key=item_dict["key"],
                    value=item_dict["value"],
                    metadata=item_dict.get("metadata", {}),
                    score=1.0,
                    created_at=datetime.now(timezone.utc),
                ))
        return items

    async def delete(self, key: str) -> bool:
        redis_key = self._make_key(key)
        return bool(await self._redis.delete(redis_key))

    async def clear(self, prefix: str = "") -> int:
        """清除指定前缀的所有 Working Memory"""
        pattern = self._make_key(f"{prefix}*")
        count = 0
        async for key in self._redis.scan_iter(match=pattern):
            await self._redis.delete(key)
            count += 1
        return count