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

"""Context processing module.

Handles processing of subtask contexts including:
- Attachments (text documents and images for vision models)
- Knowledge bases (RAG retrieval)
- Selected documents (direct injection for notebook mode)

Replaces the original attachments.py with unified context support.

This module provides unified context processing based on user_subtask_id,
eliminating the need to pass separate attachment_ids and knowledge_base_ids.
"""

import logging
from typing import Any, List, Optional, Tuple

from fastapi import HTTPException, status
from langchain_core.tools import BaseTool
from sqlalchemy import or_, select, update
from sqlalchemy.orm import Session

from app.models.knowledge import KnowledgeDocument
from app.models.subtask import Subtask
from app.models.subtask_context import ContextStatus, ContextType, SubtaskContext
from app.services.context import context_service
from shared.models.db import ContextStatus as DBContextStatus
from shared.models.knowledge import (
    ChatContextsResult,
    KnowledgeBaseToolAccessMode,
    KnowledgeBaseToolsResult,
)
from shared.prompts import (
    KB_PROMPT_RELAXED,
    KB_PROMPT_RESTRICTED_ANALYST,
    KB_PROMPT_STRICT,
)

logger = logging.getLogger(__name__)

# Table context prompt template - will be dynamically generated with table info
TABLE_PROMPT_TEMPLATE = """

<table_context>
# IMPORTANT: Data Table Context - HIGHEST PRIORITY

The user has selected data table(s) for this conversation. This indicates that the user's request is related to these tables.

## Available Tables:
{tables_info}

## Critical Rules - READ CAREFULLY:
1. **ALWAYS assume the user's request is about the selected table(s)** unless they explicitly mention otherwise
2. When the user says "分析" (analyze), "查看" (view), "统计" (statistics), or similar words, they mean to analyze/view the selected table(s)
3. **DO NOT** use other analysis tools (like PaaS analysis) when a table is selected - use the table tool instead
4. **You MUST pass `provider`, `base_id`, and `sheet_id_or_name` parameters** when calling the table tool

## Available Table Tool:
- `data_table_query`: Query table data including both schema (field definitions) and records (data rows)

## Workflow:
1. Call `data_table_query` with the correct `provider`, `base_id`, and `sheet_id_or_name` to get table schema and data
2. Analyze the returned data based on user's request
3. Present the results

The user explicitly selected these table(s) - prioritize table operations over any other tools.
</table_context>
"""


def build_table_prompt(table_contexts: List[dict]) -> str:
    """
    Build dynamic TABLE_PROMPT with actual table information.

    Args:
        table_contexts: List of table context dicts with name, provider, baseId, sheetIdOrName

    Returns:
        Formatted prompt string with table info
    """
    if not table_contexts:
        return ""

    tables_info_lines = []
    for i, ctx in enumerate(table_contexts, 1):
        name = ctx.get("name", f"Table {i}")
        provider = ctx.get("provider", "")
        base_id = ctx.get("baseId", "")
        sheet_id = ctx.get("sheetIdOrName", "")
        tables_info_lines.append(
            f"{i}. **{name}**\n"
            f"   - provider: `{provider}`\n"
            f"   - base_id: `{base_id}`\n"
            f"   - sheet_id_or_name: `{sheet_id}`"
        )

    tables_info = "\n".join(tables_info_lines)
    return TABLE_PROMPT_TEMPLATE.format(tables_info=tables_info)


async def process_contexts(
    db: Session,
    context_ids: List[int],
    message: str,
) -> str | list[dict[str, Any]]:
    """
    Process multiple contexts and build message with all context contents.

    Args:
        db: Database session (SQLAlchemy Session)
        context_ids: List of context IDs
        message: Original message

    Returns:
        Message with all context contents prepended, or OpenAI Responses API
        format vision content list for images
    """
    if not context_ids:
        return message

    # Collect all contexts
    text_contents = []
    image_contents = []

    for idx, context_id in enumerate(context_ids, start=1):
        try:
            context = context_service.get_context_optional(
                db=db,
                context_id=context_id,
            )

            if context is None:
                logger.warning(f"Context {context_id} not found")
                continue

            if context.status != ContextStatus.READY.value:
                logger.warning(f"Context {context_id} is not ready: {context.status}")
                continue

            # Process based on context type
            if context.context_type == ContextType.ATTACHMENT.value:
                _process_attachment_context(context, idx, text_contents, image_contents)
            elif context.context_type == ContextType.KNOWLEDGE_BASE.value:
                # Knowledge base contexts are handled via RAG tools, not here
                logger.debug(
                    f"Knowledge base context {context_id} will be used via RAG"
                )

        except (ValueError, KeyError) as e:
            logger.exception(f"Error processing context {context_id}")
            continue
        except Exception as e:
            logger.exception(f"Unexpected error processing context {context_id}")
            continue

    # Build vision structure if images present
    if image_contents:
        return _build_vision_structure(text_contents, image_contents, message)

    # Combine text contents if present
    if text_contents:
        return _combine_text_contents(text_contents, message)

    return message


def _build_vision_structure(
    text_contents: List[str],
    image_contents: List[dict],
    message: str,
) -> list[dict[str, Any]]:
    """
    Build OpenAI Responses API format vision content for image contexts.

    The returned block order is:
      1. [optional] <attachment> text block — metadata headers for images and
         extracted text from any accompanying document attachments.
      2. One ``input_image`` block per image.
      3. The user's own message as a standalone ``input_text`` block.

    Separating the user message into its own block keeps it stable across turns
    (good for prefix caching) and makes the intent unambiguous to the model.

    Args:
        text_contents: List of text content strings (attachment contents without XML tags)
        image_contents: List of image content dictionaries
        message: User message

    Returns:
        List of content blocks in OpenAI Responses API format
    """
    content: list[dict[str, Any]] = []

    # Collect all attachment content parts (text documents and image headers)
    all_attachment_parts: list[str] = []

    # Add text attachment contents
    if text_contents:
        all_attachment_parts.extend(text_contents)

    # Add image metadata headers (inside the attachment tag)
    for img in image_contents:
        if "image_header" in img:
            all_attachment_parts.append(img["image_header"])

    # 1. Attachment metadata block (only when there is something to show)
    if all_attachment_parts:
        attachment_text = (
            "<attachment>" + "".join(all_attachment_parts) + "</attachment>"
        )
        content.append({"type": "input_text", "text": attachment_text})

    # 2. Image blocks
    for img in image_contents:
        image_base64 = img.get("image_base64", "")
        mime_type = img.get("mime_type", "image/jpeg")
        if image_base64:
            content.append(
                {
                    "type": "input_image",
                    "image_url": f"data:{mime_type};base64,{image_base64}",
                }
            )

    # 3. User message as its own text block — keeps it isolated from attachment
    #    metadata so the model sees the question without extra noise, and so the
    #    exact user text is preserved for prefix-cache stability.
    content.append({"type": "input_text", "text": message})

    return content


def _combine_text_contents(
    text_contents: List[str], message: str
) -> list[dict[str, Any]]:
    """Combine text contents with user message.

    Returns a list of content blocks in OpenAI Responses API format so that
    the downstream converter can cleanly separate the user message from
    system context (attachment metadata) and wrap them in ``<system-reminder>``.

    Args:
        text_contents: List of text content strings (attachment contents without XML tags)
        message: Original user message

    Returns:
        List of content blocks: attachment metadata block(s) followed by the
        user message block (always last).
    """
    attachment_text = "<attachment>" + "".join(text_contents) + "</attachment>"
    return [
        {"type": "input_text", "text": attachment_text},
        {"type": "input_text", "text": message},
    ]


def _process_attachment_context(
    context: SubtaskContext,
    idx: int,
    text_contents: List[str],
    image_contents: List[dict],
    task_id: Optional[int] = None,
    subtask_id: Optional[int] = None,
) -> None:
    """
    Process an attachment context and add to appropriate list.

    Args:
        context: The SubtaskContext record
        idx: Attachment index (for labeling)
        text_contents: List to append text content to
        image_contents: List to append image content to
        task_id: Optional task ID for building sandbox path
        subtask_id: Optional subtask ID for building sandbox path
    """
    # Check if it's an image attachment
    if context_service.is_image_context(context) and context.image_base64:
        # Build image attachment metadata
        attachment_id = context.id
        filename = context.original_filename
        mime_type = context.mime_type or "unknown"
        file_size = context.file_size or 0
        formatted_size = context_service.format_file_size(file_size)
        url = context_service.build_attachment_url(attachment_id)

        # Build sandbox path if task_id and subtask_id are provided
        sandbox_path = context_service.build_sandbox_path(task_id, subtask_id, filename)

        # Build image metadata header with optional sandbox path
        if sandbox_path:
            image_header = (
                f"[Image Attachment: {filename} | ID: {attachment_id} | "
                f"Type: {mime_type} | Size: {formatted_size} | URL: {url} | "
                f"File Path(already in sandbox): {sandbox_path}]"
            )
        else:
            image_header = (
                f"[Image Attachment: {filename} | ID: {attachment_id} | "
                f"Type: {mime_type} | Size: {formatted_size} | URL: {url}]"
            )

        image_contents.append(
            {
                "image_base64": context.image_base64,
                "mime_type": context.mime_type,
                "filename": context.original_filename,
                "id": attachment_id,
                "url": url,
                "image_header": image_header,
            }
        )
    else:
        # Text document - get formatted content with attachment index
        # The content is wrapped in <attachment> XML tags by context_service
        doc_prefix = context_service.build_document_text_prefix(
            context,
            task_id=task_id,
            subtask_id=subtask_id,
        )
        if doc_prefix:
            text_contents.append(f"[Attachment {idx}]\n{doc_prefix}")


async def process_attachments(
    db: Any,
    attachment_ids: List[int],
    user_id: int,
    message: str,
) -> str | list[dict[str, Any]]:
    """
    Process multiple attachments and build message with all attachment contents.

    This is a backward-compatible wrapper around process_contexts.

    Args:
        db: Database session (SQLAlchemy Session)
        attachment_ids: List of attachment IDs (now context IDs)
        user_id: User ID (unused, kept for backward compatibility with callers)
        message: Original message

    Returns:
        Message with all attachment contents prepended, or OpenAI Responses API
        format vision content list for images
    """
    return await process_contexts(db, attachment_ids, message)


def extract_knowledge_base_ids(
    db: Session,
    context_ids: List[int],
) -> List[int]:
    """
    Extract knowledge base IDs from context list.

    Args:
        db: Database session
        context_ids: List of context IDs

    Returns:
        List of knowledge_id values from knowledge_base type contexts
    """
    if not context_ids:
        return []

    contexts = (
        db.query(SubtaskContext)
        .filter(
            SubtaskContext.id.in_(context_ids),
            SubtaskContext.context_type == ContextType.KNOWLEDGE_BASE.value,
        )
        .all()
    )

    return [c.knowledge_id for c in contexts if c.knowledge_id]


def _validate_attachment_ownership(
    db: Session,
    attachment_ids: List[int],
    user_id: int,
    task_id: Optional[int] = None,
) -> List[int]:
    """Validate attachment ownership and state.

    Validates that:
    - Attachments belong to the current user
    - Attachments are of type ATTACHMENT
    - Attachments are in READY status
    - Attachments are either unlinked or already linked to the same task

    This prevents attachments from being "stolen" from other active conversations
    while still allowing retry scenarios within the same task.

    Args:
        db: Database session
        attachment_ids: List of attachment IDs to validate
        user_id: User ID for ownership check
        task_id: Optional task ID for cross-subtask validation

    Returns:
        List of valid attachment IDs

    Raises:
        HTTPException: 403 if any attachment IDs are invalid or unauthorized
    """
    # Build base filters
    filters = [
        SubtaskContext.id.in_(attachment_ids),
        SubtaskContext.user_id == user_id,
        SubtaskContext.context_type == ContextType.ATTACHMENT.value,
        SubtaskContext.status == ContextStatus.READY.value,
    ]

    # Add cross-subtask validation if task_id is provided
    if task_id:
        task_subtask_ids = (
            select(Subtask.id)
            .filter(Subtask.task_id == task_id, Subtask.user_id == user_id)
            .scalar_subquery()
        )
        filters.append(
            or_(
                SubtaskContext.subtask_id == 0,  # Unlinked
                SubtaskContext.subtask_id.in_(task_subtask_ids),  # Same task
            )
        )
    else:
        # Without task_id, only allow unlinked attachments
        filters.append(SubtaskContext.subtask_id == 0)

    # Query with row locking
    valid_rows = db.query(SubtaskContext.id).filter(*filters).with_for_update().all()
    valid_ids = [row[0] for row in valid_rows]

    # Check for invalid IDs
    invalid_ids = set(attachment_ids) - set(valid_ids)
    if invalid_ids:
        raise HTTPException(
            status_code=status.HTTP_403_FORBIDDEN,
            detail=f"Invalid or unauthorized attachment IDs: {sorted(invalid_ids)}",
        )

    return valid_ids


def link_contexts_to_subtask(
    db: Session,
    subtask_id: int,
    user_id: int,
    attachment_ids: List[int] | None = None,
    contexts: List[Any] | None = None,
    task: Optional["TaskResource"] = None,
    user_name: Optional[str] = None,
) -> List[int]:
    """
    Link attachments and create knowledge base/table contexts for a subtask.

    This function handles three types of contexts in a single database transaction:
    1. Attachments: Pre-uploaded files with existing context IDs, batch update subtask_id
    2. Knowledge bases: Selected at send time, batch create SubtaskContext records
       (without extracted_text - RAG retrieval is done later via tools/Service)
    3. Tables: Selected at send time, batch create SubtaskContext records
       (table context is used for MCP tool injection)

    When knowledge bases are created, they are automatically synced to the task-level
    knowledgeBaseRefs for future use across all subtasks.

    SECURITY NOTE: When attachment_ids is provided, ownership validation is ALWAYS
    performed to prevent attachment hijacking across users/tasks.

    Args:
        db: Database session
        subtask_id: Subtask ID to link contexts to
        user_id: User ID
        attachment_ids: List of pre-uploaded attachment context IDs to link
        contexts: List of ContextItem objects from payload (for knowledge bases and tables)
        task: Optional pre-queried TaskResource object for syncing KB to task level
        user_name: Optional pre-queried user name for KB sync boundBy field

    Returns:
        List of all linked/created context IDs
    """
    # Import TaskResource for type hint
    from app.models.task import TaskResource

    linked_context_ids = []

    # Step 1: Validate and collect attachment IDs
    if attachment_ids:
        valid_attachment_ids = _validate_attachment_ownership(
            db=db,
            attachment_ids=attachment_ids,
            user_id=user_id,
            task_id=task.id if task else None,
        )
        linked_context_ids.extend(valid_attachment_ids)
    else:
        valid_attachment_ids = None

    # Prepare knowledge base, table, and selected_documents contexts for batch creation
    (
        kb_contexts_to_create,
        table_contexts_to_create,
        selected_docs_contexts_to_create,
    ) = _prepare_contexts_for_creation(contexts, subtask_id, user_id)

    # Combine all contexts to create
    all_contexts_to_create = (
        kb_contexts_to_create
        + table_contexts_to_create
        + selected_docs_contexts_to_create
    )

    # Execute all database operations in a single transaction
    try:
        created_context_ids = _batch_update_and_insert_contexts(
            db=db,
            attachment_ids=valid_attachment_ids,
            contexts_to_create=all_contexts_to_create,
            subtask_id=subtask_id,
            task_id=task.id if task else None,
        )
        linked_context_ids.extend(created_context_ids)

        # Sync subtask-level knowledge bases to task level
        if task and kb_contexts_to_create and user_name:
            _sync_kb_contexts_to_task(
                db, kb_contexts_to_create, task, user_id, user_name
            )

        # Sync attachments to running sandbox (if sandbox is healthy)
        if valid_attachment_ids and task:
            _schedule_attachment_sync_to_sandbox(
                db=db,
                task_id=task.id,
                subtask_id=subtask_id,
                attachment_ids=valid_attachment_ids,
            )

    except Exception as e:
        db.rollback()
        logger.exception(f"Failed to link contexts to subtask {subtask_id}: {e}")
        raise

    return linked_context_ids


def _sync_kb_contexts_to_task(
    db: Session,
    kb_contexts: List[SubtaskContext],
    task: "TaskResource",
    user_id: int,
    user_name: str,
) -> None:
    """
    Sync subtask-level knowledge base contexts to task-level knowledgeBaseRefs.

    This function syncs each KB selected in the subtask to the task level using
    append mode with deduplication. Failures are logged but do not raise exceptions.

    Args:
        db: Database session
        kb_contexts: List of KB SubtaskContext objects that were just created
        task: Pre-queried TaskResource object to sync KBs to
        user_id: User ID who selected the KBs
        user_name: Pre-queried user name for boundBy field
    """
    from app.services.knowledge import TaskKnowledgeBaseService

    task_kb_service = TaskKnowledgeBaseService()
    for kb_context in kb_contexts:
        knowledge_id = (
            kb_context.type_data.get("knowledge_id") if kb_context.type_data else None
        )
        if not knowledge_id:
            continue

        try:
            synced = task_kb_service.sync_subtask_kb_to_task(
                db=db,
                task=task,
                knowledge_id=knowledge_id,
                user_id=user_id,
                user_name=user_name,
            )
            if synced:
                logger.info(
                    f"[_sync_kb_contexts_to_task] Synced KB {knowledge_id} "
                    f"from subtask to task {task.id}"
                )
            else:
                logger.info(
                    f"[_sync_kb_contexts_to_task] Sync skipped for KB {knowledge_id} "
                    f"on task {task.id} selected by user {user_id}; "
                    "see sync_subtask_kb_to_task info logs for the skip reason"
                )
        except Exception as e:
            # Log but don't fail - syncing to task level is best-effort
            logger.warning(
                f"[_sync_kb_contexts_to_task] Failed to sync KB {knowledge_id} "
                f"to task {task.id}: {e}"
            )


def _schedule_attachment_sync_to_sandbox(
    db: Session,
    task_id: int,
    subtask_id: int,
    attachment_ids: List[int],
) -> None:
    """
    Schedule asynchronous sync of attachments to running sandbox.

    This function checks if the sandbox is running and healthy, and if so,
    uploads attachments to the sandbox filesystem. Failures are logged but
    do not block the main flow.

    Args:
        db: Database session
        task_id: Task ID (sandbox_id is derived from task_id)
        subtask_id: Subtask ID
        attachment_ids: List of attachment context IDs to sync
    """
    import asyncio

    from app.services.sandbox_file_syncer import sync_attachment_to_sandbox_background

    try:
        # Import sandbox sync utilities (deferred to avoid circular imports)
        # Health check is performed asynchronously in the background task
        # to avoid blocking the main flow

        # Get attachment data from database
        attachment_contexts = (
            db.query(SubtaskContext).filter(SubtaskContext.id.in_(attachment_ids)).all()
        )

        if not attachment_contexts:
            logger.debug(
                f"[_schedule_attachment_sync_to_sandbox] No attachments found "
                f"for ids={attachment_ids}"
            )
            return

        # Schedule async sync for each attachment
        for ctx in attachment_contexts:
            filename = ctx.original_filename
            if not filename:
                continue

            # Get binary data from storage backend
            binary_data = context_service.get_attachment_binary_data(db=db, context=ctx)
            if binary_data is None:
                logger.warning(
                    f"[_schedule_attachment_sync_to_sandbox] Failed to get binary data "
                    f"for attachment {ctx.id}"
                )
                continue

            # Schedule background sync
            try:
                loop = asyncio.get_running_loop()
                loop.create_task(
                    sync_attachment_to_sandbox_background(
                        task_id=task_id,
                        subtask_id=subtask_id,
                        filename=filename,
                        binary_data=binary_data,
                    )
                )
                logger.debug(
                    f"[_schedule_attachment_sync_to_sandbox] Scheduled sync for "
                    f"attachment {ctx.id} (filename={filename})"
                )
            except RuntimeError:
                # No event loop running - try main event loop
                try:
                    from app.services.chat.webpage_ws_chat_emitter import (
                        get_main_event_loop,
                    )

                    main_loop = get_main_event_loop()
                    if main_loop and main_loop.is_running():
                        asyncio.run_coroutine_threadsafe(
                            sync_attachment_to_sandbox_background(
                                task_id=task_id,
                                subtask_id=subtask_id,
                                filename=filename,
                                binary_data=binary_data,
                            ),
                            main_loop,
                        )
                        logger.debug(
                            f"[_schedule_attachment_sync_to_sandbox] Scheduled sync on "
                            f"main loop for attachment {ctx.id}"
                        )
                    else:
                        logger.debug(
                            f"[_schedule_attachment_sync_to_sandbox] No running event loop, "
                            f"skipping sync for attachment {ctx.id}"
                        )
                except Exception as e:
                    logger.debug(
                        f"[_schedule_attachment_sync_to_sandbox] Failed to schedule sync: {e}"
                    )

    except Exception as e:
        # Log but don't fail - syncing to sandbox is best-effort
        logger.warning(
            f"[_schedule_attachment_sync_to_sandbox] Error scheduling sync for "
            f"task {task_id}: {e}"
        )


def _prepare_contexts_for_creation(
    contexts: List[Any] | None,
    subtask_id: int,
    user_id: int,
) -> Tuple[List[SubtaskContext], List[SubtaskContext], List[SubtaskContext]]:
    """
    Prepare knowledge base, table, and selected_documents contexts for batch creation.

    Args:
        contexts: List of ContextItem objects from payload
        subtask_id: Subtask ID to link contexts to
        user_id: User ID

    Returns:
        Tuple of (kb_contexts, table_contexts, selected_docs_contexts) ready for insertion
    """
    kb_contexts_to_create: List[SubtaskContext] = []
    table_contexts_to_create: List[SubtaskContext] = []
    selected_docs_contexts_to_create: List[SubtaskContext] = []

    if not contexts:
        return (
            kb_contexts_to_create,
            table_contexts_to_create,
            selected_docs_contexts_to_create,
        )

    for ctx in contexts:
        if ctx.type == "knowledge_base":
            try:
                kb_data = ctx.data
                knowledge_id = kb_data.get("knowledge_id")
                kb_name = kb_data.get("name", f"Knowledge Base {knowledge_id}")
                document_count = kb_data.get("document_count")
                # Get document_ids if user referenced specific documents
                document_ids = kb_data.get("document_ids", [])

                # Build type_data
                type_data_dict = {
                    "knowledge_id": int(knowledge_id) if knowledge_id else 0,
                    "document_count": document_count,
                }
                # Only add document_ids if provided
                if document_ids:
                    type_data_dict["document_ids"] = document_ids

                # Create SubtaskContext object (not yet committed)
                kb_context = SubtaskContext(
                    subtask_id=subtask_id,
                    user_id=user_id,
                    context_type=ContextType.KNOWLEDGE_BASE.value,
                    name=kb_name,
                    status=ContextStatus.READY.value,
                    type_data=type_data_dict,
                )
                kb_contexts_to_create.append(kb_context)
            except Exception as e:
                logger.warning(f"Failed to prepare knowledge base context: {e}")
                continue

        elif ctx.type == "table":
            try:
                table_data = ctx.data
                document_id = table_data.get("document_id")
                table_name = table_data.get("name", f"Table {document_id}")
                # URL is in source_config.url from frontend TableContext
                source_config = table_data.get("source_config", {})
                table_url = source_config.get("url", "") if source_config else ""

                # Create SubtaskContext object for table (not yet committed)
                table_context = SubtaskContext(
                    subtask_id=subtask_id,
                    user_id=user_id,
                    context_type=ContextType.TABLE.value,
                    name=table_name,
                    status=ContextStatus.READY.value,
                    type_data={
                        "document_id": int(document_id) if document_id else 0,
                        "url": table_url,
                    },
                )
                table_contexts_to_create.append(table_context)
            except Exception as e:
                logger.warning(f"Failed to prepare table context: {e}")
                continue

        elif ctx.type == "selected_documents":
            # Handle selected_documents context type for notebook mode
            # This stores the selected document IDs from the DocumentPanel
            try:
                docs_data = ctx.data
                knowledge_base_id = docs_data.get("knowledge_base_id")
                document_ids = docs_data.get("document_ids", [])

                if knowledge_base_id and document_ids:
                    # Create SubtaskContext object for selected documents
                    selected_docs_context = SubtaskContext(
                        subtask_id=subtask_id,
                        user_id=user_id,
                        context_type=ContextType.SELECTED_DOCUMENTS.value,
                        name=f"Selected Documents ({len(document_ids)} files)",
                        status=ContextStatus.READY.value,
                        type_data={
                            "knowledge_base_id": int(knowledge_base_id),
                            "document_ids": document_ids,
                        },
                    )
                    selected_docs_contexts_to_create.append(selected_docs_context)
                    logger.info(
                        f"Prepared selected_documents context: kb_id={knowledge_base_id}, "
                        f"doc_ids={document_ids}"
                    )
            except Exception as e:
                logger.warning(f"Failed to prepare selected_documents context: {e}")
                continue

    return (
        kb_contexts_to_create,
        table_contexts_to_create,
        selected_docs_contexts_to_create,
    )


def _batch_update_and_insert_contexts(
    db: Session,
    attachment_ids: List[int] | None,
    contexts_to_create: List[SubtaskContext],
    subtask_id: int,
    task_id: Optional[int] = None,
) -> List[int]:
    """
    Execute batch update and insert operations for contexts.

    Args:
        db: Database session
        attachment_ids: List of validated attachment context IDs to update
        contexts_to_create: List of contexts (KB or table) to insert
        subtask_id: Subtask ID for logging
        task_id: Optional task ID for additional validation in update

    Returns:
        List of created context IDs

    Raises:
        HTTPException: 409 if attachment update count mismatch (concurrent modification)
    """
    created_context_ids = []

    # Batch update existing attachments' subtask_id with conditional validation
    if attachment_ids:
        # Build update conditions
        update_filters = [SubtaskContext.id.in_(attachment_ids)]

        # Add task-level validation if task_id is provided
        if task_id:
            task_subtask_ids = (
                select(Subtask.id).filter(Subtask.task_id == task_id).scalar_subquery()
            )
            update_filters.append(
                or_(
                    SubtaskContext.subtask_id == 0,
                    SubtaskContext.subtask_id.in_(task_subtask_ids),
                )
            )

        # Execute conditional update
        update_stmt = (
            update(SubtaskContext).where(*update_filters).values(subtask_id=subtask_id)
        )
        result = db.execute(update_stmt)

        # Verify all expected rows were updated
        if result.rowcount != len(attachment_ids):
            raise HTTPException(
                status_code=status.HTTP_409_CONFLICT,
                detail="Some attachments could not be linked. Please retry.",
            )

    # Batch add new contexts (KB and table)
    if contexts_to_create:
        db.add_all(contexts_to_create)

    # Single commit for all operations
    db.commit()

    # Refresh contexts to get their IDs
    for ctx in contexts_to_create:
        db.refresh(ctx)
        created_context_ids.append(ctx.id)
        logger.debug(
            f"Created {ctx.context_type} context: id={ctx.id}, "
            f"name={ctx.name}, subtask_id={subtask_id}, "
            f"type_data={ctx.type_data}"
        )

    # Log summary
    if attachment_ids:
        logger.info(
            f"Linked {len(attachment_ids)} attachment contexts to subtask {subtask_id}"
        )
    if contexts_to_create:
        kb_count = sum(
            1
            for c in contexts_to_create
            if c.context_type == ContextType.KNOWLEDGE_BASE.value
        )
        table_count = sum(
            1 for c in contexts_to_create if c.context_type == ContextType.TABLE.value
        )
        selected_docs_count = sum(
            1
            for c in contexts_to_create
            if c.context_type == ContextType.SELECTED_DOCUMENTS.value
        )
        logger.info(
            f"Created {kb_count} knowledge base contexts, {table_count} table contexts, "
            f"and {selected_docs_count} selected_documents contexts for subtask {subtask_id}"
        )

    return created_context_ids


# ==================== Unified Context Processing ====================


async def prepare_contexts_for_chat(
    db: Session,
    user_subtask_id: int,
    user_id: int,
    message: str,
    base_system_prompt: str,
    task_id: Optional[int] = None,
    context_window: Optional[int] = None,
    model_config: Optional[dict[str, Any]] = None,
) -> ChatContextsResult:
    """
    Unified context processing based on user_subtask_id.

    This function retrieves all contexts associated with a user subtask and:
    1. Processes attachment contexts - injects content into the message
    2. Processes knowledge base contexts - creates KnowledgeBaseTool for RAG
    3. Processes table contexts - injects table info into system prompt

    This eliminates the need to pass separate attachment_ids and knowledge_base_ids
    through the call chain.

    Args:
        db: Database session
        user_subtask_id: User subtask ID to get contexts from
        user_id: User ID for access control
        message: Original user message
        base_system_prompt: Base system prompt to enhance
        task_id: Optional task ID for resolving task-level bound KBs (group chat)
        context_window: Optional model context window size from Model spec.
            Used for selected_documents injection threshold calculation.
            If None, uses default value (128000).
        model_config: Optional model configuration used by restricted KB safe summary.
    Returns:
        ChatContextsResult with processed message, table info, and KB results.
    """
    from .tables import parse_table_url

    # Get all contexts for this subtask
    contexts = context_service.get_by_subtask(db, user_subtask_id) or []

    # Separate contexts by type
    attachment_contexts = [
        c
        for c in contexts
        if c.context_type == ContextType.ATTACHMENT.value
        and c.status == ContextStatus.READY.value
    ]
    kb_contexts = [
        c
        for c in contexts
        if c.context_type == ContextType.KNOWLEDGE_BASE.value
        and c.status == ContextStatus.READY.value
    ]
    table_contexts = [
        c
        for c in contexts
        if c.context_type == ContextType.TABLE.value
        and c.status == ContextStatus.READY.value
    ]
    selected_docs_contexts = [
        c
        for c in contexts
        if c.context_type == ContextType.SELECTED_DOCUMENTS.value
        and c.status == ContextStatus.READY.value
    ]

    logger.info(
        f"[prepare_contexts_for_chat] subtask={user_subtask_id}: "
        f"{len(attachment_contexts)} attachments, {len(kb_contexts)} knowledge bases, "
        f"{len(table_contexts)} tables, {len(selected_docs_contexts)} selected_documents"
    )

    # 1. Process attachment contexts - inject into message
    final_message = await _process_attachment_contexts_for_message(
        attachment_contexts,
        message,
        task_id=task_id,
        subtask_id=user_subtask_id,
    )

    # 2. Process knowledge base contexts - create tools
    kb_result = _prepare_kb_tools_from_contexts(
        kb_contexts=kb_contexts,
        user_id=user_id,
        db=db,
        base_system_prompt=base_system_prompt,
        task_id=task_id,
        user_subtask_id=user_subtask_id,
        model_config=model_config,
    )

    extra_tools = kb_result.extra_tools
    enhanced_system_prompt = kb_result.enhanced_system_prompt
    kb_meta_prompt = kb_result.kb_meta_prompt

    # 3. Process table contexts - create DataTableTool and build dynamic prompt
    parsed_tables = []
    if table_contexts:
        for table_ctx in table_contexts:
            logger.info(
                f"[prepare_contexts_for_chat] Processing table context: "
                f"id={table_ctx.id}, name={table_ctx.name}, type_data={table_ctx.type_data}"
            )
            table_url = (
                table_ctx.type_data.get("url", "") if table_ctx.type_data else ""
            )

            if table_url:
                table_info = parse_table_url(table_url)
                if table_info:
                    # Add table name to the parsed info
                    table_info["name"] = table_ctx.name
                    parsed_tables.append(table_info)
                    logger.info(
                        f"[prepare_contexts_for_chat] Table parsed: name={table_ctx.name}, "
                        f"baseId={table_info.get('baseId')}, "
                        f"sheetIdOrName={table_info.get('sheetIdOrName')}"
                    )
                else:
                    logger.warning(
                        f"[prepare_contexts_for_chat] Failed to parse table URL: {table_url}"
                    )
            else:
                logger.warning(
                    f"[prepare_contexts_for_chat] Table context has no URL in type_data"
                )

        # Note: DataTableTool creation is handled by chat_shell service in HTTP mode.
        # In non-HTTP mode (deprecated), the tool would be created here, but since
        # we're standardizing on HTTP mode, we only return parsed_tables for chat_shell.
        # chat_shell will create the DataTableTool when it receives table_contexts.
        if parsed_tables:
            table_prompt = build_table_prompt(parsed_tables)
            enhanced_system_prompt = f"{enhanced_system_prompt}{table_prompt}"
            logger.info(
                f"[prepare_contexts_for_chat] Added {len(parsed_tables)} table(s) to system prompt. "
                f"Table contexts will be passed to chat_shell for DataTableTool creation."
            )

    # 4. Process selected_documents contexts - direct injection or RAG fallback
    logger.info(
        f"[prepare_contexts_for_chat] SELECTED_DOCS_CHECK: subtask={user_subtask_id}, "
        f"selected_docs_contexts_count={len(selected_docs_contexts)}, "
        f"contexts_data={[c.type_data for c in selected_docs_contexts]}"
    )
    if selected_docs_contexts:
        from .selected_documents import process_selected_documents_contexts

        # Build kwargs for process_selected_documents_contexts
        # Only pass context_window if explicitly provided (from Model spec)
        kwargs = {
            "db": db,
            "selected_docs_contexts": selected_docs_contexts,
            "user_id": user_id,
            "message": final_message,
            "base_system_prompt": enhanced_system_prompt,
            "extra_tools": extra_tools,
            "user_subtask_id": user_subtask_id,
        }
        if context_window is not None:
            kwargs["context_window"] = context_window
            logger.info(
                f"[prepare_contexts_for_chat] Using model context_window={context_window} "
                f"for selected_documents injection threshold"
            )

        final_message, enhanced_system_prompt, extra_tools = (
            process_selected_documents_contexts(**kwargs)
        )

    # Derive flag from parsed_tables (not raw table_contexts) to stay consistent
    # with the returned table payload — if parsing fails, the flag stays False.
    has_table_context = len(parsed_tables) > 0

    # Rebuild KnowledgeBaseToolsResult with potentially mutated enhanced_system_prompt
    # and extra_tools (table prompt and selected_documents processing may have modified
    # them after kb_result was computed).
    final_kb = KnowledgeBaseToolsResult(
        extra_tools=extra_tools,
        enhanced_system_prompt=enhanced_system_prompt,
        kb_meta_prompt=kb_meta_prompt,
        knowledge_base_ids=kb_result.knowledge_base_ids,
        is_user_selected_kb=kb_result.is_user_selected_kb,
        document_ids=kb_result.document_ids,
        kb_tool_access_mode=kb_result.kb_tool_access_mode,
    )
    return ChatContextsResult(
        final_message=final_message,
        has_table_context=has_table_context,
        table_contexts=parsed_tables,
        kb=final_kb,
    )


async def _process_attachment_contexts_for_message(
    attachment_contexts: List[SubtaskContext],
    message: str,
    task_id: Optional[int] = None,
    subtask_id: Optional[int] = None,
) -> str | list[dict[str, Any]]:
    """
    Process attachment contexts and build message with content.

    Args:
        attachment_contexts: List of attachment SubtaskContext records
        message: Original user message
        task_id: Optional task ID for building sandbox path
        subtask_id: Optional subtask ID for building sandbox path
    Returns:
        Message with attachment contents prepended, or OpenAI Responses API
        format vision content list for images
    """
    if not attachment_contexts:
        return message

    text_contents = []
    image_contents = []

    for idx, context in enumerate(attachment_contexts, start=1):
        try:
            _process_attachment_context(
                context,
                idx,
                text_contents,
                image_contents,
                task_id=task_id,
                subtask_id=subtask_id,
            )
        except Exception as e:
            logger.exception(f"Error processing attachment context {context.id}: {e}")
            continue

    # If we have images, return a multi-vision structure with image metadata headers
    if image_contents:
        return _build_vision_structure(text_contents, image_contents, message)

    # If only text contents, combine them as list format
    if text_contents:
        attachment_text = "<attachment>" + "".join(text_contents) + "</attachment>"
        return [
            {"type": "input_text", "text": attachment_text},
            {"type": "input_text", "text": message},
        ]

    # Return original message if no attachment contents were processed
    return message


def _get_user_kb_tool_access_mode(
    db: Session,
    user_id: int,
    knowledge_base_ids: List[int],
) -> tuple[str, str]:
    """Get KB tool exposure mode for the current user.

    Restricted Analysts in a group are downgraded to search-only KB usage.
    Direct document access remains blocked in other permission checks.

    Args:
        db: Database session
        user_id: User ID
        knowledge_base_ids: List of knowledge base IDs to check

    Returns:
        Tuple of (access_mode, reason)
    """
    from app.services.share import get_knowledge_base_tool_access_mode_by_ids

    return get_knowledge_base_tool_access_mode_by_ids(db, user_id, knowledge_base_ids)


def _prepare_kb_tools_from_contexts(
    kb_contexts: List[SubtaskContext],
    user_id: int,
    db: Session,
    base_system_prompt: str,
    task_id: Optional[int] = None,
    user_subtask_id: Optional[int] = None,
    model_config: Optional[dict[str, Any]] = None,
) -> KnowledgeBaseToolsResult:
    """Prepare knowledge base tools from context records.

    NOTE:
    - Knowledge base metadata (kb_meta_prompt) is returned separately and should be injected
      via the `dynamic_context` mechanism to keep system prompts static for better caching.
    - knowledge_base_ids, is_user_selected_kb, and document_ids are returned so callers
      can populate ExecutionRequest fields without extra DB queries.
    """

    extra_tools: List[BaseTool] = []
    enhanced_system_prompt = base_system_prompt
    kb_meta_prompt = ""

    # Priority 1: Subtask-level knowledge bases (user-selected for this message)
    subtask_kb_ids = [c.knowledge_id for c in kb_contexts if c.knowledge_id is not None]

    # Track whether KB is user-selected (strict mode) or inherited from task (relaxed mode)
    is_user_selected_kb = bool(subtask_kb_ids)

    # Determine which knowledge bases to use based on priority
    if subtask_kb_ids:
        # Use subtask-level KBs only (user's explicit selection takes precedence)
        knowledge_base_ids = subtask_kb_ids
        logger.info(
            f"[_prepare_kb_tools_from_contexts] Using {len(knowledge_base_ids)} "
            f"subtask-level knowledge bases (priority 1, strict mode): {knowledge_base_ids}"
        )
    elif task_id:
        # Priority 2: Fall back to task-level bound knowledge bases
        knowledge_base_ids = _get_bound_knowledge_base_ids(db, task_id)
        if knowledge_base_ids:
            logger.info(
                f"[_prepare_kb_tools_from_contexts] Using {len(knowledge_base_ids)} "
                f"task-level bound knowledge bases (priority 2, relaxed mode): {knowledge_base_ids}"
            )
    else:
        knowledge_base_ids = []

    # Extract document_ids from subtask KB contexts (no extra DB query needed).
    # Normalize to int, skip invalid values, and deduplicate while preserving order.
    document_ids: List[int] = []
    if is_user_selected_kb:
        seen_doc_ids: set[int] = set()
        for c in kb_contexts:
            if c.type_data and isinstance(c.type_data, dict):
                raw_doc_ids = c.type_data.get("document_ids", [])
                if isinstance(raw_doc_ids, list):
                    for doc_id in raw_doc_ids:
                        try:
                            normalized = int(doc_id)
                        except (TypeError, ValueError):
                            logger.warning(
                                "[_prepare_kb_tools_from_contexts] Ignore invalid "
                                "document_id=%s in context_id=%s",
                                doc_id,
                                getattr(c, "id", None),
                            )
                            continue
                        if normalized not in seen_doc_ids:
                            seen_doc_ids.add(normalized)
                            document_ids.append(normalized)

    if not knowledge_base_ids:
        return KnowledgeBaseToolsResult(
            extra_tools=extra_tools,
            enhanced_system_prompt=enhanced_system_prompt,
            kb_meta_prompt="",
            knowledge_base_ids=[],
            is_user_selected_kb=False,
            document_ids=[],
            kb_tool_access_mode=KnowledgeBaseToolAccessMode.FULL,
        )

    kb_tool_access_mode, access_reason = _get_user_kb_tool_access_mode(
        db, user_id, knowledge_base_ids
    )
    is_restricted_search_only = (
        kb_tool_access_mode == KnowledgeBaseToolAccessMode.RESTRICTED_SEARCH_ONLY
    )

    if is_restricted_search_only:
        logger.warning(
            "[_prepare_kb_tools_from_contexts] User %s downgraded to restricted "
            "search-only KB mode for knowledge bases %s, reason: %s",
            user_id,
            knowledge_base_ids,
            access_reason,
        )

    logger.info(
        f"[_prepare_kb_tools_from_contexts] Creating KnowledgeBaseTool for "
        f"{len(knowledge_base_ids)} knowledge bases: {knowledge_base_ids}"
    )

    # Import KnowledgeBaseTool
    from chat_shell.tools.builtin import KnowledgeBaseTool

    # Create KnowledgeBaseTool with the specified knowledge bases
    # KB configs (max_calls, exempt_calls, name) are now fetched from Backend API
    kb_tool = KnowledgeBaseTool(
        knowledge_base_ids=knowledge_base_ids,
        document_ids=document_ids or [],
        user_id=user_id,
        db_session=db,
        user_subtask_id=user_subtask_id,
        current_model_name=(model_config or {}).get("model_name"),
        current_model_namespace=(model_config or {}).get("model_namespace")
        or "default",
        injection_mode="hybrid",
        tool_access_mode=kb_tool_access_mode,
    )
    extra_tools.append(kb_tool)

    if is_restricted_search_only:
        kb_instruction = KB_PROMPT_RESTRICTED_ANALYST
        kb_meta_prompt = _build_kb_meta_prompt(db, knowledge_base_ids, restricted=True)
        logger.info(
            "[_prepare_kb_tools_from_contexts] Using RESTRICTED_ANALYST mode prompt "
            "(search-only KB access)"
        )
    else:
        # Build KB meta prompt for dynamic_context injection.
        # This is based on the resolved KB IDs of the CURRENT request (no DB scanning by task).
        kb_meta_prompt = _build_kb_meta_prompt(db, knowledge_base_ids)

        # Choose prompt template based on whether KB is user-selected or inherited from task.
        # Keep KB prompt templates fully static (no kb_meta_list placeholder).
        if is_user_selected_kb:
            # Strict mode: User explicitly selected KB for this message
            kb_instruction = KB_PROMPT_STRICT
            logger.info(
                "[_prepare_kb_tools_from_contexts] Using STRICT mode prompt "
                "(user explicitly selected KB)"
            )
        else:
            # Relaxed mode: KB inherited from task, AI can use general knowledge as fallback
            kb_instruction = KB_PROMPT_RELAXED
            logger.info(
                "[_prepare_kb_tools_from_contexts] Using RELAXED mode prompt "
                "(KB inherited from task)"
            )

    enhanced_system_prompt = f"{base_system_prompt}{kb_instruction}"

    return KnowledgeBaseToolsResult(
        extra_tools=extra_tools,
        enhanced_system_prompt=enhanced_system_prompt,
        kb_meta_prompt=kb_meta_prompt,
        knowledge_base_ids=knowledge_base_ids,
        is_user_selected_kb=is_user_selected_kb,
        document_ids=document_ids,
        kb_tool_access_mode=kb_tool_access_mode,
    )


def _get_bound_knowledge_base_ids(db: Session, task_id: int) -> List[int]:
    """
    Get knowledge base IDs bound to a task.

    This function delegates to task_knowledge_base_service.get_bound_knowledge_base_ids()
    for the actual lookup and migration logic. It wraps the call with additional
    error handling to ensure chat functionality is never blocked.

    The underlying service uses ID-priority lookup with name fallback for
    backward compatibility. Legacy refs (name-only) are automatically
    migrated to include the ID field.

    Args:
        db: Database session
        task_id: Task ID

    Returns:
        List of knowledge base IDs bound to the task
    """
    from app.services.knowledge.task_knowledge_base_service import (
        task_knowledge_base_service,
    )

    logger.info(f"[_get_bound_knowledge_base_ids] START task_id={task_id}")

    try:
        kb_ids = task_knowledge_base_service.get_bound_knowledge_base_ids(db, task_id)
        logger.info(
            f"[_get_bound_knowledge_base_ids] RESULT task_id={task_id}, kb_ids={kb_ids}"
        )
        return kb_ids
    except Exception as e:
        # Catch all exceptions to ensure robustness - this function should
        # never block chat functionality even if KB lookup fails
        logger.warning(f"Failed to get bound KB IDs for task {task_id}: {e}")
        return []


def _build_kb_meta_prompt(
    db: Session,
    knowledge_base_ids: List[int],
    restricted: bool = False,
) -> str:
    """Build KB meta prompt for dynamic_context injection.

    IMPORTANT:
    - The formatter lives in `kb_meta.py` and MUST NOT query DB.
    - This function assembles KB meta from resolved KB IDs of the current request.

    Args:
        db: Database session
        knowledge_base_ids: Resolved KB IDs for the current request

    Returns:
        Formatted prompt string, or empty string.
    """

    if not knowledge_base_ids:
        return ""

    try:
        from app.services.chat.preprocessing.kb_meta import (
            format_kb_meta_prompt,
            format_restricted_kb_meta_prompt,
            select_kb_summary_text,
        )
        from app.services.knowledge import KnowledgeService
        from app.services.knowledge.task_knowledge_base_service import (
            task_knowledge_base_service,
        )

        kb_map = task_knowledge_base_service.get_knowledge_bases_by_ids(
            db, knowledge_base_ids
        )
        kb_prompt_stats = KnowledgeService.get_document_prompt_stats(
            db, knowledge_base_ids
        )

        kb_meta_list: list[dict[str, Any]] = []
        for kb_id in knowledge_base_ids:
            kb_kind = kb_map.get(kb_id)
            if not kb_kind:
                kb_meta_list.append(
                    {
                        "kb_id": kb_id,
                        "kb_name": "Unknown",
                        "search_available": False,
                        "total_document_count": 0,
                        "searchable_document_count": 0,
                        "spreadsheet_document_count": 0,
                        "summary_text": "",
                        "topics": [],
                    }
                )
                continue

            kb_spec = kb_kind.json.get("spec", {}) if kb_kind.json else {}
            # Prefer spec-level name, then model-level Kind.name, then "Unknown".
            kb_name = kb_spec.get("name") or getattr(kb_kind, "name", None) or "Unknown"

            summary_text = ""
            topics: list[str] = []
            try:
                summary_data = kb_spec.get("summary", {})
                if (
                    kb_spec.get("summaryEnabled")
                    and summary_data.get("status") == "completed"
                ):
                    summary_text = select_kb_summary_text(
                        summary_data, len(knowledge_base_ids)
                    )
                    topics = summary_data.get("topics", []) or []
            except Exception as e:
                logger.warning(
                    "[kb_meta] Failed to extract summary for KB %s: %s",
                    kb_id,
                    e,
                    exc_info=True,
                )

            stats = kb_prompt_stats.get(
                kb_id,
                {
                    "total_document_count": 0,
                    "searchable_document_count": 0,
                    "spreadsheet_document_count": 0,
                },
            )
            retrieval_config = kb_spec.get("retrievalConfig") or {}

            kb_meta_list.append(
                {
                    "kb_id": kb_id,
                    "kb_name": kb_name,
                    "search_available": bool(retrieval_config.get("retriever_name")),
                    "total_document_count": stats["total_document_count"],
                    "searchable_document_count": stats["searchable_document_count"],
                    "spreadsheet_document_count": stats["spreadsheet_document_count"],
                    "summary_text": summary_text,
                    "topics": topics,
                }
            )

        if restricted:
            return format_restricted_kb_meta_prompt(kb_meta_list)
        return format_kb_meta_prompt(kb_meta_list)
    except Exception as e:
        logger.warning(
            "Failed to build KB meta prompt for knowledge_base_ids=%s: %s",
            knowledge_base_ids,
            e,
            exc_info=True,
        )
        return ""


def get_knowledge_base_ids_from_subtask(
    db: Session,
    subtask_id: int,
) -> List[int]:
    """
    Get knowledge base IDs from a subtask's contexts.

    This is a convenience function to extract KB IDs from subtask contexts
    without needing to pass them through the call chain.

    Args:
        db: Database session
        subtask_id: Subtask ID

    Returns:
        List of knowledge_id values from knowledge_base type contexts
    """
    kb_contexts = context_service.get_knowledge_base_contexts_by_subtask(db, subtask_id)
    return [c.knowledge_id for c in kb_contexts if c.knowledge_id is not None]


def get_document_ids_from_subtask(
    db: Session,
    subtask_id: int,
) -> List[int]:
    """
    Get document IDs from a subtask's knowledge base contexts.

    When a user references specific documents from a knowledge base,
    the document_ids are stored in the context's type_data field.
    This function extracts all document IDs from all KB contexts.

    Args:
        db: Database session
        subtask_id: Subtask ID

    Returns:
        List of document_id values from knowledge_base type contexts
    """
    kb_contexts = context_service.get_knowledge_base_contexts_by_subtask(db, subtask_id)
    document_ids = []
    for c in kb_contexts:
        if c.type_data and isinstance(c.type_data, dict):
            doc_ids = c.type_data.get("document_ids", [])
            if doc_ids:
                document_ids.extend(doc_ids)
    return document_ids


def get_attachment_context_ids_from_subtask(
    db: Session,
    subtask_id: int,
) -> List[int]:
    """
    Get attachment context IDs from a subtask.

    Args:
        db: Database session
        subtask_id: Subtask ID

    Returns:
        List of attachment context IDs
    """
    attachments = context_service.get_attachments_by_subtask(db, subtask_id)
    return [a.id for a in attachments]


def get_table_context_ids_from_subtask(
    db: Session,
    subtask_id: int,
) -> List[int]:
    """
    Get table context IDs from a subtask.

    Args:
        db: Database session
        subtask_id: Subtask ID

    Returns:
        List of table context IDs
    """
    contexts = context_service.get_by_subtask(db, subtask_id)
    return [
        c.id
        for c in contexts
        if c.context_type == ContextType.TABLE.value
        and c.status == ContextStatus.READY.value
    ]
