RAGStack-Lambda

"""
Sync Status Checker Lambda

Periodically checks the status of documents waiting for KB sync (SYNC_QUEUED)
and updates their status to INDEXED or INGESTION_FAILED based on actual
Bedrock KB document status.

This Lambda is triggered by EventBridge on a schedule (e.g., every 2 minutes).

Flow:
1. Query tracking table for documents with status=SYNC_QUEUED
2. For each document, check its status in Bedrock KB
3. Update tracking table:
   - INDEXED in KB → status = INDEXED
   - FAILED in KB → status = INGESTION_FAILED
   - Still processing → leave as SYNC_QUEUED
"""

import logging
import os
from datetime import UTC, datetime

import boto3
from botocore.exceptions import ClientError

from ragstack_common.appsync import publish_image_update
from ragstack_common.config import get_config_manager_or_none, get_knowledge_base_config
from ragstack_common.ingestion import batch_check_document_statuses

logger = logging.getLogger()
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))

# Initialize AWS clients
dynamodb = boto3.resource("dynamodb")

# Maximum documents to process per invocation (to stay within Lambda timeout)
MAX_DOCUMENTS_PER_RUN = 100


def get_sync_queued_documents(table_name: str) -> list[dict]:
    """
    Query tracking table for documents with SYNC_QUEUED status.

    Returns list of document dicts with document_id, output_s3_uri, etc.
    """
    table = dynamodb.Table(table_name)

    # Scan for SYNC_QUEUED status (could use GSI for efficiency at scale)
    # Note: For production with many documents, consider adding a GSI on status
    documents = []

    try:
        response = table.scan(
            FilterExpression="#status = :status",
            ExpressionAttributeNames={"#status": "status"},
            ExpressionAttributeValues={":status": "SYNC_QUEUED"},
            Limit=MAX_DOCUMENTS_PER_RUN,
        )
        documents.extend(response.get("Items", []))

        # Handle pagination if needed
        while "LastEvaluatedKey" in response and len(documents) < MAX_DOCUMENTS_PER_RUN:
            response = table.scan(
                FilterExpression="#status = :status",
                ExpressionAttributeNames={"#status": "status"},
                ExpressionAttributeValues={":status": "SYNC_QUEUED"},
                ExclusiveStartKey=response["LastEvaluatedKey"],
                Limit=MAX_DOCUMENTS_PER_RUN - len(documents),
            )
            documents.extend(response.get("Items", []))

    except ClientError as e:
        logger.error(f"Error scanning tracking table: {e}")

    return documents[:MAX_DOCUMENTS_PER_RUN]


def update_document_status(
    table_name: str,
    document_id: str,
    new_status: str,
    error_message: str | None = None,
) -> None:
    """Update a single document's status in tracking table."""
    table = dynamodb.Table(table_name)
    now = datetime.now(UTC).isoformat()

    try:
        update_expr = "SET #status = :status, updated_at = :updated_at"
        expr_names = {"#status": "status"}
        expr_values = {
            ":status": new_status,
            ":updated_at": now,
        }

        if error_message:
            update_expr += ", error_message = :error"
            expr_values[":error"] = error_message
            # Remove ingestion_job_id on completion (success or failure)
            update_expr += " REMOVE ingestion_job_id"
        else:
            # Remove both ingestion_job_id and any stale error_message on success
            update_expr += " REMOVE ingestion_job_id, error_message"

        table.update_item(
            Key={"document_id": document_id},
            UpdateExpression=update_expr,
            ExpressionAttributeNames=expr_names,
            ExpressionAttributeValues=expr_values,
        )
        logger.info(f"Updated {document_id} status to {new_status}")

    except ClientError as e:
        logger.error(f"Failed to update status for {document_id}: {e}")


def lambda_handler(event, context):
    """
    Check status of SYNC_QUEUED documents and update accordingly.

    Triggered by EventBridge schedule rule.
    """
    # Get KB config from DynamoDB (with env var fallback)
    config_manager = get_config_manager_or_none()
    kb_id, ds_id = get_knowledge_base_config(config_manager)
    logger.info(f"Using KB config: kb_id={kb_id}, ds_id={ds_id}")

    tracking_table = os.environ.get("TRACKING_TABLE")
    graphql_endpoint = os.environ.get("GRAPHQL_ENDPOINT")

    if not tracking_table:
        raise ValueError("TRACKING_TABLE is required")

    # Get documents waiting for sync
    documents = get_sync_queued_documents(tracking_table)

    if not documents:
        logger.info("No documents with SYNC_QUEUED status")
        return {"checked": 0, "updated": 0}

    logger.info(f"Checking status for {len(documents)} SYNC_QUEUED documents")

    # Build mapping of S3 URI to document info
    uri_to_doc = {}
    for doc in documents:
        # For images, check both the image file and caption file
        # The caption.txt is what gets indexed for text search
        caption_uri = doc.get("caption_s3_uri")
        output_uri = doc.get("output_s3_uri")

        # Use caption_uri if available (image), otherwise output_uri (document)
        check_uri = caption_uri or output_uri
        if check_uri:
            uri_to_doc[check_uri] = doc

    if not uri_to_doc:
        logger.warning("No valid S3 URIs found for SYNC_QUEUED documents")
        return {"checked": 0, "updated": 0}

    # Batch check document statuses in Bedrock KB
    s3_uris = list(uri_to_doc.keys())
    statuses = batch_check_document_statuses(kb_id, ds_id, s3_uris)

    # Process results and update tracking table
    updated_count = 0
    indexed_count = 0
    failed_count = 0

    for uri, kb_status in statuses.items():
        doc = uri_to_doc.get(uri)
        if not doc:
            continue

        document_id = doc.get("document_id")
        filename = doc.get("filename", "unknown")
        doc_type = doc.get("type", "document")

        if kb_status == "INDEXED":
            # Document successfully indexed
            update_document_status(tracking_table, document_id, "INDEXED")
            updated_count += 1
            indexed_count += 1

            # Publish real-time update for images
            if doc_type == "image" and graphql_endpoint:
                try:
                    publish_image_update(
                        graphql_endpoint,
                        document_id,
                        filename,
                        "INDEXED",
                    )
                except Exception as e:
                    logger.warning(f"Failed to publish update for {document_id}: {e}")

        elif kb_status == "FAILED":
            # Document ingestion failed
            update_document_status(
                tracking_table,
                document_id,
                "INGESTION_FAILED",
                "KB ingestion job reported FAILED status",
            )
            updated_count += 1
            failed_count += 1

            # Publish real-time update for images
            if doc_type == "image" and graphql_endpoint:
                try:
                    publish_image_update(
                        graphql_endpoint,
                        document_id,
                        filename,
                        "INGESTION_FAILED",
                        error_message="KB sync failed",
                    )
                except Exception as e:
                    logger.warning(f"Failed to publish update for {document_id}: {e}")

        else:
            # Still processing (STARTING, IN_PROGRESS) or UNKNOWN
            logger.debug(f"Document {document_id} status: {kb_status}, leaving as SYNC_QUEUED")

    logger.info(
        f"Status check complete: {len(documents)} checked, "
        f"{updated_count} updated ({indexed_count} indexed, {failed_count} failed)"
    )

    return {
        "checked": len(documents),
        "updated": updated_count,
        "indexed": indexed_count,
        "failed": failed_count,
    }