# SPDX-FileCopyrightText: 2025 Weibo, Inc.
#
# SPDX-License-Identifier: Apache-2.0

"""
Session manager for Chat Shell.

Manages chat history and session state in Redis for multi-turn conversations.
Also manages cancellation state for streaming chat requests using Redis
for cross-worker communication in multi-worker deployments.

Additionally manages streaming content and blocks for mixed content rendering,
supporting page refresh recovery during streaming.
"""

import asyncio
import json
import logging
from datetime import datetime
from typing import Any, Dict, List, Optional

from app.core.cache import cache_manager
from app.core.config import settings
from shared.models.blocks import BlockStatus, create_text_block, create_tool_block

logger = logging.getLogger(__name__)

# Redis key prefix for cancellation flags
CANCEL_KEY_PREFIX = "chat:cancel:"
# Cancellation flag TTL in seconds (5 minutes should be enough for any chat)
CANCEL_FLAG_TTL = 300

# Redis key prefix for streaming content cache
STREAMING_KEY_PREFIX = "chat:streaming:"
# Redis Pub/Sub channel prefix for streaming updates
STREAMING_CHANNEL_PREFIX = "chat:stream_channel:"
# Redis key prefix for task-level streaming status (for group chat)
TASK_STREAMING_KEY_PREFIX = "chat:task_streaming:"
# Redis Pub/Sub channel prefix for callback-based SSE streaming (ClaudeCode/Agno/Dify)
CALLBACK_CHANNEL_PREFIX = "callback:channel:"
# Unified TTL for all streaming-related data (1 hour)
STREAMING_TTL = 3600


class SessionManager:
    """
    Manages chat session state in Redis.

    Stores conversation history for multi-turn chat support.
    Uses task_id as the session identifier.

    Also manages cancellation state for streaming chat requests.
    Uses Redis for cancellation flags to support multi-worker deployments.
    Uses subtask_id as the cancellation identifier.
    """

    def __init__(self):
        self._cache = cache_manager
        # Local asyncio events for in-process signaling (optimization)
        # Key: subtask_id, Value: asyncio.Event
        self._local_events: Dict[int, asyncio.Event] = {}

    def _get_history_key(self, task_id: int) -> str:
        """Generate Redis key for chat history."""
        return f"chat:history:{task_id}"

    async def get_chat_history(self, task_id: int) -> List[Dict[str, str]]:
        """
        Get chat history for a task.

        Args:
            task_id: The task ID to get history for

        Returns:
            List of message dictionaries with 'role' and 'content' keys
        """
        try:
            key = self._get_history_key(task_id)
            history = await self._cache.get(key)

            if history is None:
                return []

            # Ensure we return a list
            if isinstance(history, list):
                return history

            logger.warning(
                f"Invalid history format for task {task_id}, returning empty list"
            )
            return []

        except Exception as e:
            logger.error(f"Error getting chat history for task {task_id}: {e}")
            return []

    async def save_chat_history(
        self, task_id: int, messages: List[Dict[str, str]], expire: Optional[int] = None
    ) -> bool:
        """
        Save chat history for a task.

        Args:
            task_id: The task ID to save history for
            messages: List of message dictionaries
            expire: Optional expiration time in seconds

        Returns:
            bool: True if save was successful
        """
        try:
            key = self._get_history_key(task_id)

            # Limit history size to prevent token overflow
            max_messages = settings.CHAT_HISTORY_MAX_MESSAGES
            if len(messages) > max_messages:
                messages = messages[-max_messages:]
                logger.info(
                    f"Truncated chat history for task {task_id} to {max_messages} messages"
                )

            expire_time = expire or settings.CHAT_HISTORY_EXPIRE_SECONDS
            return await self._cache.set(key, messages, expire=expire_time)

        except Exception as e:
            logger.error(f"Error saving chat history for task {task_id}: {e}")
            return False

    async def append_message(self, task_id: int, role: str, content: str) -> bool:
        """
        Append a single message to chat history.

        Args:
            task_id: The task ID
            role: Message role ('user', 'assistant', or 'system')
            content: Message content

        Returns:
            bool: True if append was successful
        """
        try:
            history = await self.get_chat_history(task_id)
            history.append({"role": role, "content": content})
            return await self.save_chat_history(task_id, history)

        except Exception as e:
            logger.error(f"Error appending message for task {task_id}: {e}")
            return False

    async def append_user_and_assistant_messages(
        self, task_id: int, user_message: Any, assistant_message: str
    ) -> bool:
        """
        Append both user and assistant messages to chat history.

        This is the common pattern after a successful chat completion.

        Args:
            task_id: The task ID
            user_message: The user's message (string or vision dict)
            assistant_message: The assistant's response

        Returns:
            bool: True if append was successful
        """
        try:
            history = await self.get_chat_history(task_id)

            # Normalize user message content for storage
            # If it's a vision message dict, convert to standard OpenAI format
            if isinstance(user_message, dict) and user_message.get("type") == "vision":
                # Convert vision message to OpenAI format (array of content blocks)
                user_content = [
                    {"type": "text", "text": user_message.get("text", "")},
                    {
                        "type": "image_url",
                        "image_url": {
                            "url": f"data:{user_message['mime_type']};base64,{user_message['image_base64']}"
                        },
                    },
                ]
            elif isinstance(user_message, str):
                # Regular text message
                user_content = user_message
            else:
                # Fallback: convert to string
                user_content = str(user_message)

            history.append({"role": "user", "content": user_content})
            history.append({"role": "assistant", "content": assistant_message})
            return await self.save_chat_history(task_id, history)

        except Exception as e:
            logger.error(f"Error appending messages for task {task_id}: {e}")
            return False

    async def clear_history(self, task_id: int) -> bool:
        """
        Clear chat history for a task.

        Args:
            task_id: The task ID to clear history for

        Returns:
            bool: True if clear was successful
        """
        try:
            key = self._get_history_key(task_id)
            return await self._cache.delete(key)

        except Exception as e:
            logger.error(f"Error clearing chat history for task {task_id}: {e}")
            return False

    async def get_history_length(self, task_id: int) -> int:
        """
        Get the number of messages in chat history.

        Args:
            task_id: The task ID

        Returns:
            int: Number of messages in history
        """
        history = await self.get_chat_history(task_id)
        return len(history)

    # ==================== Cancellation Management ====================

    def _get_cancel_key(self, subtask_id: int) -> str:
        """Generate Redis key for cancellation flag."""
        return f"{CANCEL_KEY_PREFIX}{subtask_id}"

    async def register_stream(self, subtask_id: int) -> asyncio.Event:
        """
        Register a new streaming request and return its cancellation event.

        Creates a local asyncio.Event for in-process signaling and
        clears any existing cancellation flag in Redis.

        Args:
            subtask_id: The subtask ID for the stream

        Returns:
            asyncio.Event: Event that will be set when cancellation is requested
        """
        # Create local event for in-process signaling
        cancel_event = asyncio.Event()
        self._local_events[subtask_id] = cancel_event

        # Clear any existing cancellation flag in Redis (in case of retry)
        cancel_key = self._get_cancel_key(subtask_id)
        try:
            await self._cache.delete(cancel_key)
        except Exception as e:
            logger.warning(f"Failed to clear cancel flag for subtask {subtask_id}: {e}")

        return cancel_event

    async def cancel_stream(self, subtask_id: int) -> bool:
        """
        Request cancellation of a streaming request.

        Sets cancellation flag in Redis (for cross-worker communication)
        and also sets local event if the stream is in this process.

        Args:
            subtask_id: The subtask ID to cancel

        Returns:
            bool: True if cancellation flag was set successfully
        """
        cancel_key = self._get_cancel_key(subtask_id)

        # Set cancellation flag in Redis (cross-worker)
        try:
            success = await self._cache.set(cancel_key, True, expire=CANCEL_FLAG_TTL)
        except Exception as e:
            logger.error(
                f"Failed to set Redis cancel flag for subtask {subtask_id}: {e}"
            )
            success = False

        # Also set local event if stream is in this process (optimization)
        local_event = self._local_events.get(subtask_id)
        if local_event:
            local_event.set()

        return success

    async def unregister_stream(self, subtask_id: int):
        """
        Unregister a streaming request (cleanup after completion or cancellation).

        Removes local event and cleans up Redis cancellation flag.

        Args:
            subtask_id: The subtask ID to unregister
        """
        # Clean up local event
        if subtask_id in self._local_events:
            del self._local_events[subtask_id]

        # Clean up Redis cancellation flag
        cancel_key = self._get_cancel_key(subtask_id)
        try:
            await self._cache.delete(cancel_key)
        except Exception as e:
            logger.warning(
                f"Failed to delete cancel flag for subtask {subtask_id}: {e}"
            )

    async def is_cancelled(self, subtask_id: int) -> bool:
        """
        Check if a streaming request has been cancelled.

        Checks both local event (fast path) and Redis flag (cross-worker).
        If Redis flag is set, also sets local event for consistency.

        Args:
            subtask_id: The subtask ID to check

        Returns:
            bool: True if cancellation has been requested
        """
        # Fast path: check local event first
        local_event = self._local_events.get(subtask_id)
        if local_event and local_event.is_set():
            return True

        # Slow path: check Redis flag (for cross-worker cancellation)
        cancel_key = self._get_cancel_key(subtask_id)
        try:
            redis_flag = await self._cache.get(cancel_key)
            if redis_flag is True:
                # Set local event for consistency
                if local_event:
                    local_event.set()
                return True
        except Exception as e:
            logger.warning(
                f"Failed to check Redis cancel flag for subtask {subtask_id}: {e}"
            )

        return False

    # ==================== Streaming Content Cache ====================

    def _get_streaming_key(self, subtask_id: int) -> str:
        """Generate Redis key for streaming content cache."""
        return f"{STREAMING_KEY_PREFIX}{subtask_id}"

    async def save_streaming_content(
        self, subtask_id: int, content: str, expire: int = None
    ) -> bool:
        """
        Save streaming content to Redis (temporary cache).

        This is used for fast recovery when user refreshes during streaming.
        Content is saved frequently (every 1 second) to minimize data loss.

        Args:
            subtask_id: Subtask ID
            content: Current accumulated content
            expire: Expiration time in seconds (default from settings)

        Returns:
            bool: True if save was successful
        """
        try:
            key = self._get_streaming_key(subtask_id)
            expire_time = expire or settings.STREAMING_REDIS_TTL
            return await self._cache.set(key, content, expire=expire_time)
        except Exception as e:
            logger.error(
                f"Error saving streaming content for subtask {subtask_id}: {e}"
            )
            return False

    async def get_streaming_content(self, subtask_id: int) -> Optional[str]:
        """
        Get streaming content from Redis cache.

        Used for recovery when user refreshes during streaming.

        Args:
            subtask_id: Subtask ID

        Returns:
            str or None: Cached streaming content, or None if not found
        """
        try:
            key = self._get_streaming_key(subtask_id)
            # Use direct Redis client to get raw string content
            # (add_text_content uses redis APPEND which stores raw strings)
            redis_client = await self._cache._get_client()
            try:
                content = await redis_client.get(key)
                if content is not None:
                    # Content is stored as raw bytes, decode to string
                    if isinstance(content, bytes):
                        result = content.decode("utf-8")
                    else:
                        result = str(content)
                    logger.debug(
                        f"[SessionManager] get_streaming_content: subtask_id={subtask_id}, "
                        f"content_len={len(result)}, content_preview={result[:100] if result else 'None'}..."
                    )
                    return result
                logger.debug(
                    f"[SessionManager] get_streaming_content: subtask_id={subtask_id}, "
                    f"content=None (key not found)"
                )
                return None
            finally:
                await redis_client.aclose()
        except Exception as e:
            logger.error(
                f"Error getting streaming content for subtask {subtask_id}: {e}"
            )
            return None

    async def delete_streaming_content(self, subtask_id: int) -> bool:
        """
        Delete streaming content from Redis cache.

        Called after streaming completes (success, cancel, or error).

        Args:
            subtask_id: Subtask ID

        Returns:
            bool: True if delete was successful
        """
        try:
            key = self._get_streaming_key(subtask_id)
            return await self._cache.delete(key)
        except Exception as e:
            logger.error(
                f"Error deleting streaming content for subtask {subtask_id}: {e}"
            )
            return False

    def _get_channel_key(self, subtask_id: int) -> str:
        """Generate Redis Pub/Sub channel key for streaming updates."""
        return f"{STREAMING_CHANNEL_PREFIX}{subtask_id}"

    async def publish_streaming_chunk(self, subtask_id: int, chunk: str) -> bool:
        """
        Publish a streaming chunk to Redis Pub/Sub.

        This allows other clients (e.g., reconnected browsers) to receive
        real-time updates for an ongoing stream.

        Args:
            subtask_id: Subtask ID
            chunk: Content chunk to publish

        Returns:
            bool: True if publish was successful
        """
        try:
            channel = self._get_channel_key(subtask_id)
            # Get a Redis client for pub/sub
            redis_client = await self._cache._get_client()
            try:
                await redis_client.publish(channel, chunk)
                return True
            finally:
                await redis_client.aclose()
        except Exception as e:
            logger.error(
                f"Error publishing streaming chunk for subtask {subtask_id}: {e}"
            )
            return False

    async def publish_streaming_done(
        self, subtask_id: int, result: Optional[Dict[str, Any]] = None
    ) -> bool:
        """
        Publish a "done" signal to Redis Pub/Sub with optional result data.

        The done signal is JSON-encoded and contains:
        - __type__: "STREAM_DONE" (marker to identify this as a done signal)
        - result: The final result data (optional)

        Args:
            subtask_id: Subtask ID
            result: Optional result data to include in the done signal

        Returns:
            bool: True if publish was successful
        """
        try:
            channel = self._get_channel_key(subtask_id)
            redis_client = await self._cache._get_client()
            try:
                # Encode done signal with result data
                done_message = json.dumps({"__type__": "STREAM_DONE", "result": result})
                await redis_client.publish(channel, done_message)
                return True
            finally:
                await redis_client.aclose()
        except Exception as e:
            logger.error(f"Error publishing stream done for subtask {subtask_id}: {e}")
            return False

    async def subscribe_streaming_channel(self, subtask_id: int):
        """
        Subscribe to a streaming channel for real-time updates.

        Args:
            subtask_id: Subtask ID

        Returns:
            Tuple of (Redis client, PubSub object) or (None, None)
            Caller is responsible for closing the client when done.
        """
        try:
            channel = self._get_channel_key(subtask_id)
            redis_client = await self._cache._get_client()
            pubsub = redis_client.pubsub()
            await pubsub.subscribe(channel)
            # Return both client and pubsub so caller can close client when done
            return redis_client, pubsub
        except Exception as e:
            logger.error(
                f"Error subscribing to streaming channel for subtask {subtask_id}: {e}"
            )
            return None, None

    # ==================== Callback Event Pub/Sub (ClaudeCode/Agno streaming) ====================

    def _get_callback_channel_key(self, subtask_id: int) -> str:
        """Generate Redis Pub/Sub channel key for callback-based streaming."""
        return f"{CALLBACK_CHANNEL_PREFIX}{subtask_id}"

    async def publish_callback_event(self, subtask_id: int, event: Any) -> bool:
        """Publish an execution event to the callback stream channel.

        Used by the /internal/callback handler to forward executor events to
        any SSE consumers that are streaming a ClaudeCode/Agno/Dify task.

        Args:
            subtask_id: Subtask ID
            event: ExecutionEvent instance with to_dict() method

        Returns:
            bool: True if publish was successful
        """
        try:
            channel = self._get_callback_channel_key(subtask_id)
            event_json = json.dumps(event.to_dict())
            redis_client = await self._cache._get_client()
            try:
                await redis_client.publish(channel, event_json)
                return True
            finally:
                await redis_client.aclose()
        except Exception as e:
            logger.error(
                f"[SessionManager] publish_callback_event failed for subtask {subtask_id}: {e}"
            )
            return False

    async def subscribe_callback_channel(self, subtask_id: int):
        """Subscribe to the callback event channel for a subtask.

        Returns (redis_client, pubsub) so the caller can poll with
        pubsub.get_message() and close the client when done.

        Args:
            subtask_id: Subtask ID

        Returns:
            Tuple of (redis_client, pubsub) or (None, None) on failure.
        """
        try:
            channel = self._get_callback_channel_key(subtask_id)
            redis_client = await self._cache._get_client()
            pubsub = redis_client.pubsub()
            await pubsub.subscribe(channel)
            return redis_client, pubsub
        except Exception as e:
            logger.error(
                f"[SessionManager] subscribe_callback_channel failed for subtask {subtask_id}: {e}"
            )
            return None, None

    # ==================== Task-Level Streaming Status (Group Chat) ====================

    def _get_task_streaming_key(self, task_id: int) -> str:
        """Generate Redis key for task-level streaming status."""
        return f"{TASK_STREAMING_KEY_PREFIX}{task_id}"

    async def set_task_streaming_status(
        self,
        task_id: int,
        subtask_id: int,
        user_id: int,
        username: str,
    ) -> bool:
        """
        Set task-level streaming status (used for group chat).

        Args:
            task_id: Task ID
            subtask_id: Subtask ID that is streaming
            user_id: User who triggered the stream
            username: Username of the user

        Returns:
            bool: True if set was successful
        """
        try:
            key = self._get_task_streaming_key(task_id)
            now_iso = datetime.now().isoformat()
            value = {
                "subtask_id": subtask_id,
                "user_id": user_id,
                "username": username,
                "started_at": now_iso,
                "last_activity_at": now_iso,
            }
            logger.info(
                f"[SessionManager] set_task_streaming_status: key={key}, "
                f"task_id={task_id}, subtask_id={subtask_id}, user_id={user_id}, "
                f"expire={STREAMING_TTL}"
            )
            result = await self._cache.set(key, value, expire=STREAMING_TTL)
            logger.info(f"[SessionManager] set_task_streaming_status result: {result}")
            return result
        except Exception as e:
            logger.error(
                f"Error setting task streaming status for task {task_id}: {e}",
                exc_info=True,
            )
            return False

    async def touch_task_streaming_activity(self, task_id: int) -> bool:
        """
        Refresh task-level streaming last activity timestamp.

        Args:
            task_id: Task ID

        Returns:
            bool: True if updated successfully, False otherwise
        """
        try:
            key = self._get_task_streaming_key(task_id)
            status = await self._cache.get(key)
            if not status:
                return False

            if not isinstance(status, dict):
                logger.warning(
                    "[SessionManager] Invalid task streaming status format for task_id=%s",
                    task_id,
                )
                return False

            status["last_activity_at"] = datetime.now().isoformat()
            return await self._cache.set(key, status, expire=STREAMING_TTL)
        except Exception as e:
            logger.error(
                f"Error touching task streaming activity for task {task_id}: {e}",
                exc_info=True,
            )
            return False

    async def get_task_streaming_status(self, task_id: int) -> Optional[Dict[str, Any]]:
        """
        Get task-level streaming status.

        Args:
            task_id: Task ID

        Returns:
            dict or None: Streaming status data, or None if not streaming
        """
        try:
            key = self._get_task_streaming_key(task_id)
            logger.info(
                f"[SessionManager] get_task_streaming_status: key={key}, task_id={task_id}"
            )
            result = await self._cache.get(key)
            logger.info(f"[SessionManager] get_task_streaming_status result: {result}")
            return result
        except Exception as e:
            logger.error(
                f"Error getting task streaming status for task {task_id}: {e}",
                exc_info=True,
            )
            return None

    async def clear_task_streaming_status(self, task_id: int) -> bool:
        """
        Clear task-level streaming status.

        Args:
            task_id: Task ID

        Returns:
            bool: True if clear was successful
        """
        try:
            key = self._get_task_streaming_key(task_id)
            logger.info(
                f"[SessionManager] clear_task_streaming_status: key={key}, task_id={task_id}"
            )
            result = await self._cache.delete(key)
            logger.info(
                f"[SessionManager] clear_task_streaming_status result: {result}"
            )
            return result
        except Exception as e:
            logger.error(
                f"Error clearing task streaming status for task {task_id}: {e}",
                exc_info=True,
            )
            return False

    # ==================== Blocks Collection (Mixed Content Rendering) ====================
    # Uses Redis atomic operations for high performance:
    # - accumulated_content: APPEND for O(1) content addition
    # - blocks: RPUSH for O(1) block addition
    # - current_text_block_id: separate key for O(1) read/write

    def _get_blocks_key(self, subtask_id: int) -> str:
        """Generate Redis key for blocks list."""
        return f"{STREAMING_KEY_PREFIX}blocks:{subtask_id}"

    def _get_current_text_block_key(self, subtask_id: int) -> str:
        """Generate Redis key for current text block ID."""
        return f"{STREAMING_KEY_PREFIX}text_block:{subtask_id}"

    async def add_tool_block(
        self,
        subtask_id: int,
        tool_use_id: str,
        tool_name: str,
        tool_input: Optional[Dict[str, Any]] = None,
        display_name: Optional[str] = None,
        tool_protocol: Optional[str] = None,
        server_label: Optional[str] = None,
    ) -> None:
        """Add a tool block for a subtask.

        This also finalizes any current text block before upserting the tool block.
        Reuses block ids to make duplicate start callbacks idempotent.

        Args:
            subtask_id: Subtask ID
            tool_use_id: Tool use ID
            tool_name: Tool name
            tool_input: Tool input parameters
            display_name: Optional display name for the tool
            tool_protocol: Optional Responses protocol type
            server_label: Optional MCP server label
        """
        try:
            # Finalize current text block before adding tool block
            await self._finalize_current_text_block(subtask_id)

            # Create tool block using unified function
            block = create_tool_block(
                tool_use_id=tool_use_id,
                tool_name=tool_name,
                tool_input=tool_input,
                display_name=display_name,
                tool_protocol=tool_protocol,
                server_label=server_label,
            )

            # Upsert by block id so callback retries do not duplicate tool blocks.
            await self.add_block(subtask_id, block)

            logger.info(
                f"[SessionManager] Upserted tool block for subtask {subtask_id}: "
                f"id={block['id']}, tool_name={tool_name}"
            )
        except Exception as e:
            logger.error(
                f"[SessionManager] Failed to add tool block for subtask {subtask_id}: {e}"
            )

    async def update_tool_block_status(
        self,
        subtask_id: int,
        tool_use_id: str,
        status: Optional[str] = None,
        tool_output: Optional[str] = None,
        tool_input: Optional[Dict[str, Any]] = None,
        tool_protocol: Optional[str] = None,
        server_label: Optional[str] = None,
    ) -> None:
        """Update tool block status, output, and/or input.

        This is a convenience wrapper around add_block() for tool blocks.
        It retrieves the existing block, updates the specified fields, and
        calls add_block() with upsert semantics.

        Args:
            subtask_id: Subtask ID
            tool_use_id: Tool use ID
            status: New status (optional, e.g. "done", "error")
            tool_output: Optional tool output to set
            tool_input: Optional tool input/arguments to update (used by interactive_form_question MCP tool)
            tool_protocol: Optional Responses protocol type
            server_label: Optional MCP server label
        """
        try:
            # Get existing blocks to find the tool block
            blocks = await self.get_blocks(subtask_id)
            existing_block = None
            for block in blocks:
                if (
                    block.get("type") == "tool"
                    and block.get("tool_use_id") == tool_use_id
                ):
                    existing_block = block
                    break

            if existing_block:
                # Update existing block
                if status is not None:
                    existing_block["status"] = status
                if tool_output is not None:
                    existing_block["tool_output"] = tool_output
                if tool_input is not None:
                    existing_block["tool_input"] = tool_input
                if tool_protocol is not None:
                    existing_block["tool_protocol"] = tool_protocol
                if server_label is not None:
                    existing_block["server_label"] = server_label
                await self.add_block(subtask_id, existing_block)
                logger.debug(
                    f"[SessionManager] Updated tool block for subtask {subtask_id}: "
                    f"id={tool_use_id}, status={status}, "
                    f"has_tool_input={tool_input is not None}"
                )
            else:
                logger.warning(
                    f"[SessionManager] Tool block not found for subtask {subtask_id}: "
                    f"tool_use_id={tool_use_id}"
                )
        except Exception as e:
            logger.error(
                f"[SessionManager] Failed to update tool block status for subtask {subtask_id}: {e}"
            )

    async def add_block(self, subtask_id: int, block: Dict[str, Any]) -> None:
        """Add or update a block for the subtask.

        This is the unified method for block management (upsert semantics):
        - If block with same id exists -> update it
        - If block doesn't exist -> append it
        - If no id provided -> generate one and append

        Used for custom block types like subscription_preview, video, image,
        and also as the internal implementation for tool block updates.

        Args:
            subtask_id: Subtask ID
            block: Block data dict with at least 'type', optionally 'id'
        """
        try:
            # Finalize current text block before adding new block (to maintain order)
            await self._finalize_current_text_block(subtask_id)

            # Ensure block has id
            block_id = block.get("id")
            if not block_id:
                block_id = f"block-{int(asyncio.get_event_loop().time() * 1000)}"
                block["id"] = block_id

            # Ensure block has timestamp
            if "timestamp" not in block:
                block["timestamp"] = int(asyncio.get_event_loop().time() * 1000)

            blocks_key = self._get_blocks_key(subtask_id)
            redis_client = await self._cache._get_client()
            try:
                # Check if block with this id already exists (upsert logic)
                blocks_raw = await redis_client.lrange(blocks_key, 0, -1)
                existing_index = None
                for i, block_json in enumerate(blocks_raw):
                    existing_block = json.loads(block_json)
                    if existing_block.get("id") == block_id:
                        existing_index = i
                        break

                if existing_index is not None:
                    # Update existing block
                    await redis_client.lset(
                        blocks_key, existing_index, json.dumps(block)
                    )
                    logger.debug(
                        f"[SessionManager] Updated block for subtask {subtask_id}: "
                        f"id={block_id}, type={block.get('type')}"
                    )
                else:
                    # Append new block
                    await redis_client.rpush(blocks_key, json.dumps(block))
                    logger.info(
                        f"[SessionManager] Added block for subtask {subtask_id}: "
                        f"id={block_id}, type={block.get('type')}"
                    )

                await redis_client.expire(blocks_key, STREAMING_TTL)
            finally:
                await redis_client.aclose()

        except Exception as e:
            logger.error(
                f"[SessionManager] Failed to add block for subtask {subtask_id}: {e}"
            )

    async def add_text_content(self, subtask_id: int, content: str) -> None:
        """Add text content to the current text block.

        Creates a new text block if there isn't one currently active.
        Uses Redis APPEND for O(1) content addition.

        Args:
            subtask_id: Subtask ID
            content: Text content to add
        """
        if not content:
            return

        try:
            # Append to accumulated content using Redis APPEND (O(1))
            streaming_key = self._get_streaming_key(subtask_id)
            text_block_key = self._get_current_text_block_key(subtask_id)
            blocks_key = self._get_blocks_key(subtask_id)

            logger.debug(
                f"[SessionManager] add_text_content: subtask_id={subtask_id}, "
                f"content_len={len(content)}, streaming_key={streaming_key}"
            )

            redis_client = await self._cache._get_client()
            try:
                # Append content to streaming cache
                new_len = await redis_client.append(streaming_key, content)
                await redis_client.expire(streaming_key, STREAMING_TTL)
                logger.debug(
                    f"[SessionManager] add_text_content: appended to Redis, "
                    f"subtask_id={subtask_id}, new_total_len={new_len}"
                )

                # Check if we have a current text block
                current_block_id = await redis_client.get(text_block_key)

                if current_block_id:
                    # Update existing text block's content
                    # We need to find and update the last text block
                    blocks_raw = await redis_client.lrange(blocks_key, -1, -1)
                    if blocks_raw:
                        last_block = json.loads(blocks_raw[0])
                        if (
                            last_block.get("id") == current_block_id.decode()
                            if isinstance(current_block_id, bytes)
                            else current_block_id
                        ):
                            last_block["content"] = (
                                last_block.get("content", "") + content
                            )
                            await redis_client.lset(
                                blocks_key, -1, json.dumps(last_block)
                            )
                else:
                    # Create new text block
                    block = create_text_block(content=content)
                    await redis_client.rpush(blocks_key, json.dumps(block))
                    await redis_client.set(
                        text_block_key, block["id"], ex=STREAMING_TTL
                    )
                    await redis_client.expire(blocks_key, STREAMING_TTL)
                    logger.info(
                        f"[SessionManager] Created text block for subtask {subtask_id}: "
                        f"id={block['id']}"
                    )
            finally:
                await redis_client.aclose()
        except Exception as e:
            logger.error(
                f"[SessionManager] Failed to add text content for subtask {subtask_id}: {e}"
            )

    async def _finalize_current_text_block(self, subtask_id: int) -> None:
        """Finalize the current text block by setting status to done."""
        try:
            text_block_key = self._get_current_text_block_key(subtask_id)
            blocks_key = self._get_blocks_key(subtask_id)

            redis_client = await self._cache._get_client()
            try:
                current_block_id = await redis_client.get(text_block_key)
                if not current_block_id:
                    return

                block_id = (
                    current_block_id.decode()
                    if isinstance(current_block_id, bytes)
                    else current_block_id
                )

                # Find and update the text block
                blocks_raw = await redis_client.lrange(blocks_key, 0, -1)
                for i, block_json in enumerate(blocks_raw):
                    block = json.loads(block_json)
                    if block.get("id") == block_id:
                        block["status"] = BlockStatus.DONE.value
                        await redis_client.lset(blocks_key, i, json.dumps(block))
                        break

                # Clear current text block ID
                await redis_client.delete(text_block_key)
            finally:
                await redis_client.aclose()
        except Exception as e:
            logger.warning(
                f"[SessionManager] Failed to finalize text block for subtask {subtask_id}: {e}"
            )

    async def get_blocks(self, subtask_id: int) -> List[Dict[str, Any]]:
        """Get all blocks for a subtask without finalizing.

        This is used for page refresh recovery to get the current state
        of blocks during streaming without modifying them.

        Args:
            subtask_id: Subtask ID

        Returns:
            List of all blocks for the subtask
        """
        try:
            blocks_key = self._get_blocks_key(subtask_id)
            redis_client = await self._cache._get_client()
            try:
                blocks_raw = await redis_client.lrange(blocks_key, 0, -1)
                blocks = [json.loads(b) for b in blocks_raw] if blocks_raw else []
                logger.debug(
                    f"[SessionManager] get_blocks for subtask {subtask_id}: "
                    f"count={len(blocks)}"
                )
                return blocks
            finally:
                await redis_client.aclose()
        except Exception as e:
            logger.error(
                f"[SessionManager] Failed to get blocks for subtask {subtask_id}: {e}"
            )
            return []

    async def finalize_and_get_blocks(self, subtask_id: int) -> List[Dict[str, Any]]:
        """Finalize any pending text block and return all blocks.

        This should be called when the subtask completes (DONE event).

        Args:
            subtask_id: Subtask ID

        Returns:
            List of all blocks for the subtask
        """
        try:
            # Finalize current text block
            await self._finalize_current_text_block(subtask_id)

            # Get all blocks
            blocks_key = self._get_blocks_key(subtask_id)
            redis_client = await self._cache._get_client()
            try:
                blocks_raw = await redis_client.lrange(blocks_key, 0, -1)
                blocks = [json.loads(b) for b in blocks_raw] if blocks_raw else []
                logger.info(
                    f"[SessionManager] Finalized blocks for subtask {subtask_id}: "
                    f"count={len(blocks)}"
                )
                return blocks
            finally:
                await redis_client.aclose()
        except Exception as e:
            logger.error(
                f"[SessionManager] Failed to get blocks for subtask {subtask_id}: {e}"
            )
            return []

    async def get_accumulated_content(self, subtask_id: int) -> str:
        """Get accumulated content for a subtask.

        Args:
            subtask_id: Subtask ID

        Returns:
            Accumulated content string
        """
        # Use the existing streaming content cache
        content = await self.get_streaming_content(subtask_id)
        return content or ""

    async def cleanup_streaming_state(
        self, subtask_id: int, task_id: Optional[int] = None
    ) -> None:
        """Clean up all streaming state for a completed subtask.

        Args:
            subtask_id: Subtask ID
            task_id: Optional Task ID for clearing task-level streaming status
        """
        try:
            streaming_key = self._get_streaming_key(subtask_id)
            blocks_key = self._get_blocks_key(subtask_id)
            text_block_key = self._get_current_text_block_key(subtask_id)

            redis_client = await self._cache._get_client()
            try:
                await redis_client.delete(streaming_key, blocks_key, text_block_key)
                logger.debug(
                    f"[SessionManager] Cleaned up streaming state for subtask {subtask_id}"
                )
            finally:
                await redis_client.aclose()

            # Also clear task-level streaming status if task_id is provided
            if task_id:
                await self.clear_task_streaming_status(task_id)
        except Exception as e:
            logger.warning(
                f"[SessionManager] Failed to cleanup streaming state for subtask {subtask_id}: {e}"
            )


# Global session manager instance
session_manager = SessionManager()
