RAGStack-Lambda

"""
Scrape Discover Lambda

Processes URLs from the discovery queue, extracts links, and adds
new discovered URLs back to the queue (recursive discovery).

Input event (SQS triggered):
{
    "Records": [{
        "body": "{\"job_id\": \"uuid\", \"url\": \"https://...\", \"depth\": 0}"
    }]
}

Output:
{
    "processed": 1,
    "discovered": 5,
    "skipped": 2
}
"""

import json
import logging
import os
import time
from datetime import UTC, datetime

import boto3
from botocore.exceptions import ClientError

from ragstack_common.appsync import publish_scrape_update
from ragstack_common.scraper import ScrapePage, ScrapeStatus, UrlStatus
from ragstack_common.scraper.discovery import (
    extract_links,
    filter_discovered_urls,
    normalize_url,
)
from ragstack_common.scraper.fetcher import HttpFetcher
from ragstack_common.scraper.models import ScrapeConfig

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


def lambda_handler(event, context):
    """
    Main Lambda handler - processes discovery queue messages.
    """
    # Get environment variables
    jobs_table = os.environ.get("SCRAPE_JOBS_TABLE")
    urls_table = os.environ.get("SCRAPE_URLS_TABLE")
    discovery_queue_url = os.environ.get("SCRAPE_DISCOVERY_QUEUE_URL")
    processing_queue_url = os.environ.get("SCRAPE_PROCESSING_QUEUE_URL")
    request_delay_ms = int(os.environ.get("REQUEST_DELAY_MS", "500"))

    if not jobs_table:
        raise ValueError("SCRAPE_JOBS_TABLE environment variable required")
    if not urls_table:
        raise ValueError("SCRAPE_URLS_TABLE environment variable required")
    if not discovery_queue_url:
        raise ValueError("SCRAPE_DISCOVERY_QUEUE_URL environment variable required")
    if not processing_queue_url:
        raise ValueError("SCRAPE_PROCESSING_QUEUE_URL environment variable required")

    dynamodb = boto3.resource("dynamodb")
    jobs_tbl = dynamodb.Table(jobs_table)
    urls_tbl = dynamodb.Table(urls_table)
    sqs = boto3.client("sqs")

    processed = 0
    discovered = 0
    skipped = 0
    batch_item_failures = []

    # Process SQS records
    for record in event.get("Records", []):
        try:
            message = json.loads(record["body"])
            job_id = message["job_id"]
            url = message["url"]
            depth = message.get("depth", 0)

            logger.info(f"Processing discovery: job={job_id}, url={url}, depth={depth}")

            # Check if job is still active
            job_response = jobs_tbl.get_item(Key={"job_id": job_id})
            job_item = job_response.get("Item")

            if not job_item:
                logger.warning(f"Job not found: {job_id}")
                continue

            job_status = job_item.get("status")
            if job_status in [
                ScrapeStatus.CANCELLED.value,
                ScrapeStatus.FAILED.value,
                ScrapeStatus.COMPLETED.value,
            ]:
                logger.info(f"Job {job_id} is {job_status}, skipping")
                skipped += 1
                continue

            # Normalize URL for deduplication
            normalized_url = normalize_url(url)

            # Check if URL already visited
            existing = urls_tbl.get_item(Key={"job_id": job_id, "url": normalized_url})
            if existing.get("Item"):
                logger.info(f"URL already visited: {normalized_url}")
                skipped += 1
                continue

            # Create page record (mark as discovered)
            page = ScrapePage(
                job_id=job_id,
                url=normalized_url,
                status=UrlStatus.PENDING,
                depth=depth,
            )

            page_data = page.to_dict()
            urls_tbl.put_item(Item=page_data)

            # Get job config
            config_data = job_item.get("config", {})
            config = ScrapeConfig.from_dict(config_data)
            base_url = job_item.get("base_url", url)

            # Fetch the page to extract links
            fetcher = HttpFetcher(
                delay_ms=request_delay_ms,
                cookies=config.cookies,
                headers=config.headers,
            )
            result = fetcher.fetch(normalized_url)

            if result.error:
                logger.warning(f"Fetch failed during discovery: {normalized_url} - {result.error}")
                # Mark URL as failed but continue - processing will retry
                urls_tbl.update_item(
                    Key={"job_id": job_id, "url": normalized_url},
                    UpdateExpression="SET #status = :status, #error = :err",
                    ExpressionAttributeNames={"#status": "status", "#error": "error"},
                    ExpressionAttributeValues={
                        ":status": UrlStatus.FAILED.value,
                        ":err": result.error,
                    },
                )
                # Update job failed count
                jobs_tbl.update_item(
                    Key={"job_id": job_id},
                    UpdateExpression="SET failed_count = failed_count + :inc, updated_at = :ts",
                    ExpressionAttributeValues={
                        ":inc": 1,
                        ":ts": datetime.now(UTC).isoformat(),
                    },
                )
                continue

            # Send URL to processing queue
            sqs.send_message(
                QueueUrl=processing_queue_url,
                MessageBody=json.dumps(
                    {
                        "job_id": job_id,
                        "url": normalized_url,
                        "depth": depth,
                    }
                ),
            )

            # Update job total URLs count
            jobs_tbl.update_item(
                Key={"job_id": job_id},
                UpdateExpression="SET total_urls = total_urls + :inc, updated_at = :ts",
                ExpressionAttributeValues={
                    ":inc": 1,
                    ":ts": datetime.now(UTC).isoformat(),
                },
            )

            # Publish discovery progress update to subscribers
            graphql_endpoint = os.environ.get("GRAPHQL_ENDPOINT")
            current_total = int(job_item.get("total_urls", 0)) + 1
            publish_scrape_update(
                graphql_endpoint=graphql_endpoint,
                job_id=job_id,
                base_url=job_item.get("base_url", url),
                title=job_item.get("title") or job_item.get("base_url", url),
                status=job_item.get("status", ScrapeStatus.DISCOVERING.value),
                total_urls=current_total,
                processed_count=int(job_item.get("processed_count", 0)),
                failed_count=int(job_item.get("failed_count", 0)),
            )

            # Extract and filter links if within depth limit
            max_depth = config.max_depth
            max_pages = config.max_pages

            if depth < max_depth and result.is_html:
                links = extract_links(result.content, normalized_url)

                # Track URLs we've already seen (including current URL)
                # DynamoDB provides cross-invocation dedup; this handles within-batch
                visited = {normalized_url}

                filtered = filter_discovered_urls(
                    urls=links,
                    base_url=base_url,
                    config=config,
                    visited=visited,
                )

                logger.info(f"Discovered {len(filtered)} new URLs from {normalized_url}")

                # Check max pages limit
                job_refresh = jobs_tbl.get_item(Key={"job_id": job_id})
                total_discovered = int(job_refresh.get("Item", {}).get("total_urls", 0))
                remaining = max_pages - total_discovered

                # Queue new URLs for discovery
                urls_to_queue = filtered[:remaining] if remaining > 0 else []

                for link in urls_to_queue:
                    sqs.send_message(
                        QueueUrl=discovery_queue_url,
                        MessageBody=json.dumps(
                            {
                                "job_id": job_id,
                                "url": link,
                                "depth": depth + 1,
                            }
                        ),
                    )
                    discovered += 1

                if remaining <= 0:
                    logger.info(f"Max pages limit ({max_pages}) reached for job {job_id}")

            processed += 1

            # Respect rate limit between pages
            time.sleep(request_delay_ms / 1000.0)

        except ClientError as e:
            error_code = e.response.get("Error", {}).get("Code", "")
            logger.error(f"AWS error processing record: {error_code} - {e}")
            batch_item_failures.append({"itemIdentifier": record["messageId"]})
        except Exception as e:
            logger.error(f"Error processing record: {e}", exc_info=True)
            batch_item_failures.append({"itemIdentifier": record["messageId"]})

    return {
        "processed": processed,
        "discovered": discovered,
        "skipped": skipped,
        "batchItemFailures": batch_item_failures,
    }