"""
索引进度推送服务

通过 Redis Pub/Sub 推送索引进度，前端通过 SSE 订阅。
"""
import json
import logging
from typing import Optional

logger = logging.getLogger(__name__)

# Redis 频道前缀
CHANNEL_PREFIX = "knowledge:indexing:progress:"


class IndexingProgressService:
    """索引进度推送服务"""

    @staticmethod
    def _channel(knowledge_base_id: str) -> str:
        return f"{CHANNEL_PREFIX}{knowledge_base_id}"

    @classmethod
    async def publish(
        cls,
        knowledge_base_id: str,
        document_id: str,
        step: str,
        progress: float,
        message: str = "",
        document_name: str = "",
        error: Optional[str] = None,
    ):
        """
        发布索引进度事件

        Args:
            knowledge_base_id: 知识库 ID
            document_id: 文档 ID
            step: 当前步骤 (extracting/cleaning/chunking/vectorizing/completed/failed)
            progress: 进度 0.0 ~ 1.0
            message: 进度描述
            document_name: 文档名称
            error: 错误信息（仅 failed 步骤）
        """
        try:
            from utils.redis import RedisClient
            client = await RedisClient.get_client()

            event = {
                "document_id": document_id,
                "document_name": document_name,
                "step": step,
                "progress": round(progress, 2),
                "message": message,
            }
            if error:
                event["error"] = error

            await client.publish(
                cls._channel(knowledge_base_id),
                json.dumps(event, ensure_ascii=False),
            )
        except Exception as e:
            logger.warning(f"发布索引进度失败: {e}")

    @classmethod
    async def subscribe(cls, knowledge_base_id: str):
        """
        订阅索引进度事件（异步生成器，用于 SSE）

        Yields:
            dict: 进度事件
        """
        from utils.redis import RedisClient
        client = await RedisClient.get_client()
        pubsub = client.pubsub()
        channel = cls._channel(knowledge_base_id)

        await pubsub.subscribe(channel)
        try:
            async for message in pubsub.listen():
                if message["type"] == "message":
                    try:
                        data = json.loads(message["data"])
                        yield data
                    except (json.JSONDecodeError, TypeError):
                        continue
        finally:
            await pubsub.unsubscribe(channel)
            await pubsub.close()