"""SharedWorkspace - Agent 间共享工作空间

基于 Redis 的共享状态存储，支持读写、订阅、锁操作。
"""

from __future__ import annotations

import json
import logging
import time
from typing import Any

logger = logging.getLogger(__name__)


class SharedWorkspace:
    """Agent 间共享工作空间

    基于 Redis 的共享状态存储，支持：
    - write/read: 读写共享数据
    - lock/unlock: 分布式锁
    - 版本控制：每次写入递增版本号
    """

    def __init__(self, redis_client: Any = None, prefix: str = "workspace"):
        """
        Args:
            redis_client: aioredis.Redis 实例，None 时使用内存字典
            prefix: Redis key 前缀
        """
        self._redis = redis_client
        self._prefix = prefix
        self._local_store: dict[str, dict[str, Any]] = {}
        self._locks: dict[str, str] = {}  # key -> lock_owner

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

    async def write(
        self, key: str, value: Any, agent_id: str, ttl: int | None = None
    ) -> int:
        """写入共享数据

        Args:
            key: 数据键
            value: 数据值
            agent_id: 写入者 ID
            ttl: 过期时间（秒），None 表示不过期

        Returns:
            版本号
        """
        entry = {
            "value": value,
            "agent_id": agent_id,
            "version": await self._get_version(key) + 1,
            "timestamp": time.time(),
        }

        if self._redis:
            redis_key = self._make_key(key)
            data = json.dumps(entry, default=str)
            if ttl:
                await self._redis.setex(redis_key, ttl, data)
            else:
                await self._redis.set(redis_key, data)
        else:
            self._local_store[key] = entry

        return entry["version"]

    async def read(self, key: str) -> dict[str, Any] | None:
        """读取共享数据

        Returns:
            {"value": ..., "agent_id": ..., "version": ..., "timestamp": ...} 或 None
        """
        if self._redis:
            redis_key = self._make_key(key)
            data = await self._redis.get(redis_key)
            if data is None:
                return None
            return json.loads(data)
        else:
            return self._local_store.get(key)

    async def delete(self, key: str) -> bool:
        """删除共享数据"""
        if self._redis:
            redis_key = self._make_key(key)
            result = await self._redis.delete(redis_key)
            return result > 0
        else:
            return self._local_store.pop(key, None) is not None

    async def lock(self, key: str, agent_id: str, timeout: float = 30.0) -> bool:
        """获取分布式锁

        Args:
            key: 要锁定的数据键
            agent_id: 请求锁的 Agent ID
            timeout: 锁超时时间（秒）

        Returns:
            是否成功获取锁
        """
        lock_key = f"{self._prefix}:lock:{key}"

        if self._redis:
            # Redis SET with NX (only if not exists) and EX (expiry)
            result = await self._redis.set(lock_key, agent_id, nx=True, ex=int(timeout))
            return result is not None
        else:
            if key in self._locks:
                return False
            self._locks[key] = agent_id
            return True

    async def unlock(self, key: str, agent_id: str) -> bool:
        """释放分布式锁

        只有锁的持有者才能释放锁。
        """
        lock_key = f"{self._prefix}:lock:{key}"

        if self._redis:
            current_owner = await self._redis.get(lock_key)
            if current_owner and current_owner.decode() == agent_id:
                await self._redis.delete(lock_key)
                return True
            return False
        else:
            if self._locks.get(key) == agent_id:
                del self._locks[key]
                return True
            return False

    async def _get_version(self, key: str) -> int:
        """获取当前版本号"""
        data = await self.read(key)
        if data is None:
            return 0
        return data.get("version", 0)

    async def list_keys(self) -> list[str]:
        """列出所有键"""
        if self._redis:
            pattern = f"{self._prefix}:*"
            keys = []
            async for key in self._redis.scan_iter(match=pattern):
                # Strip prefix
                k = key.decode() if isinstance(key, bytes) else key
                k = k[len(self._prefix) + 1:]  # Remove "prefix:"
                # Skip lock keys
                if not k.startswith("lock:"):
                    keys.append(k)
            return keys
        else:
            return list(self._local_store.keys())