"""Service Bus worker that keeps Azure AI Search in step with blob changes.

Reads Azure Event Grid blob events from an Azure Service Bus queue. For every event it checks the
blob's current state, so late, repeated or out-of-order events can't bring back old content:
- the blob exists: chunk it, upload the chunks, then delete the old chunks it no longer uses
- the blob is gone: delete all its chunks
Handles Microsoft.Storage.BlobCreated, BlobDeleted and BlobRenamed (ADLS Gen2).

extract_and_chunk_blob is a sample for the sandbox only. Replace it with your own extraction,
chunking and embedding code, then set SAMPLE_CHUNKER_IN_USE to False. While the sample is in
place, the worker starts only when USE_SAMPLE_CHUNKER=true is set.

Environment variables: SEARCH_ENDPOINT, INDEX_NAME, SERVICEBUS_FQDN, QUEUE_NAME,
USE_SAMPLE_CHUNKER (sandbox only)
Prerequisites: pip install azure-search-documents azure-servicebus azure-storage-blob azure-identity
"""
import hashlib
import json
import logging
import os
from typing import Any, Dict, List, Set

from azure.identity import DefaultAzureCredential
from azure.search.documents import SearchClient
from azure.servicebus import AutoLockRenewer, ServiceBusClient
from azure.storage.blob import BlobClient

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger("refresh-worker")

BATCH_SIZE = 1000  # Azure AI Search limit: 1,000 documents per batch
TRANSIENT_STATUS = (422, 503)  # Index Documents: 422 index temporarily unavailable, 503 service too busy

SAMPLE_CHUNKER_IN_USE = True  # set to False when you replace extract_and_chunk_blob with your own code
SAMPLE_CHUNK_SIZE = 2000  # characters per chunk in the sample chunker


class PermanentError(Exception):
    """An error that a retry won't fix. The message goes straight to the dead-letter queue."""


def normalise_url(url: str) -> str:
    # Event Grid can report the same ADLS Gen2 file with a dfs or a blob URL, so store one form
    return url.replace(".dfs.core.windows.net/", ".blob.core.windows.net/")


def blob_exists(blob_url: str, credential) -> bool:
    return BlobClient.from_blob_url(blob_url, credential=credential).exists()


def extract_and_chunk_blob(blob_url: str, credential) -> List[Dict[str, Any]]:
    """Download the blob, extract the text, chunk it and compute embeddings (project-specific).

    Sample for the sandbox only: it reads a UTF-8 .txt file, cuts the text into chunks of
    SAMPLE_CHUNK_SIZE characters and computes no embeddings. A shorter file gives fewer chunks.
    Keep chunk IDs deterministic, so a repeated event overwrites the same documents instead of
    adding new ones.
    """
    if not blob_url.lower().endswith(".txt"):
        raise PermanentError(f"The sample chunker reads only .txt files: {blob_url}")
    data = BlobClient.from_blob_url(blob_url, credential=credential).download_blob().readall()
    try:
        text = data.decode("utf-8-sig")
    except UnicodeDecodeError as exc:
        raise PermanentError(f"The sample chunker reads only UTF-8 text: {blob_url}") from exc
    doc_prefix = hashlib.sha256(blob_url.encode("utf-8")).hexdigest()[:16]
    pieces = [text[start:start + SAMPLE_CHUNK_SIZE] for start in range(0, len(text), SAMPLE_CHUNK_SIZE)]
    return [
        {"id": f"{doc_prefix}_{i}", "source_url": blob_url, "chunk_index": i, "content": piece}
        for i, piece in enumerate(pieces)
    ]


def find_chunk_keys(search_client, blob_url: str) -> Set[str]:
    safe_url = blob_url.replace("'", "''")  # escape single quotes for the OData literal
    results = search_client.search(search_text="*", filter=f"source_url eq '{safe_url}'", select=["id"])
    return {doc["id"] for doc in results if "id" in doc}


def check_results(results, action: str) -> None:
    failed = [r for r in results if not r.succeeded]
    if not failed:
        return
    first = failed[0]
    message = f"{action} failed for {len(failed)} documents, first {first.key}: {first.status_code} {first.error_message}"
    if all(r.status_code in TRANSIENT_STATUS for r in failed):
        raise RuntimeError(message)  # transient: the message is retried later
    raise PermanentError(message)


def sync_blob(search_client, credential, blob_url: str, event_time: str = "") -> None:
    """Make the index match the blob as it is now. event_time is logged to measure freshness (Step 9)."""
    old_keys = find_chunk_keys(search_client, blob_url)
    new_keys: Set[str] = set()
    if blob_exists(blob_url, credential):
        chunks = extract_and_chunk_blob(blob_url, credential)
        for start in range(0, len(chunks), BATCH_SIZE):
            check_results(search_client.merge_or_upload_documents(documents=chunks[start:start + BATCH_SIZE]), "Upload")
        new_keys = {chunk["id"] for chunk in chunks}
    # Delete only after the upload worked, so the file never disappears from search during an update
    stale = sorted(old_keys - new_keys)
    for start in range(0, len(stale), BATCH_SIZE):
        batch = [{"id": key} for key in stale[start:start + BATCH_SIZE]]
        check_results(search_client.delete_documents(documents=batch), "Delete")
    logger.info("Synced %s: %d chunks now, %d old chunks deleted (eventTime %s)",
                blob_url, len(new_keys), len(stale), event_time or "unknown")


def process_event(event: Dict[str, Any], search_client, credential) -> None:
    event_type = event.get("eventType") or event.get("type")  # Event Grid schema or CloudEvents schema
    data = event.get("data") or {}
    if event_type in ("Microsoft.Storage.BlobCreated", "Microsoft.Storage.BlobDeleted"):
        urls = [data.get("url")]
    elif event_type == "Microsoft.Storage.BlobRenamed":
        urls = [data.get("sourceUrl"), data.get("destinationUrl")]
    else:
        logger.warning("Ignoring event type: %s", event_type)
        return
    if not all(urls):
        raise PermanentError(f"Event has no blob URL: {event_type}")
    event_time = event.get("eventTime") or event.get("time") or ""  # Event Grid schema or CloudEvents schema
    for url in urls:
        sync_blob(search_client, credential, normalise_url(url), event_time)


def handle_message(message, search_client, credential) -> None:
    try:
        payload = json.loads(str(message))
    except ValueError as exc:
        raise PermanentError(f"Message body is not JSON: {exc}") from exc
    for event in payload if isinstance(payload, list) else [payload]:
        process_event(event, search_client, credential)


def run_worker() -> None:
    # Stop before reading any message, so the events stay in the queue for the real worker
    if SAMPLE_CHUNKER_IN_USE and os.environ.get("USE_SAMPLE_CHUNKER", "").lower() != "true":
        raise SystemExit("worker.py still uses the sample chunker, which is for the sandbox only. "
                         "Replace extract_and_chunk_blob with your own code and set SAMPLE_CHUNKER_IN_USE "
                         "to False, or set USE_SAMPLE_CHUNKER=true in the sandbox.")
    credential = DefaultAzureCredential()
    search_client = SearchClient(os.environ["SEARCH_ENDPOINT"], os.environ["INDEX_NAME"], credential)
    bus_client = ServiceBusClient(os.environ["SERVICEBUS_FQDN"], credential)
    queue_name = os.environ["QUEUE_NAME"]
    logger.info("Starting refresh worker on queue: %s", queue_name)

    # Iterating the receiver keeps waiting for new messages until the worker is stopped.
    # AutoLockRenewer keeps the message lock while a large file is processed.
    with bus_client, AutoLockRenewer() as renewer, \
            bus_client.get_queue_receiver(queue_name=queue_name, auto_lock_renewer=renewer) as receiver:
        for message in receiver:
            try:
                handle_message(message, search_client, credential)
                receiver.complete_message(message)
                logger.info("Successfully settled message: %s", message.message_id)
            except PermanentError as exc:
                logger.error("Dead-lettering message %s: %s", message.message_id, exc)
                receiver.dead_letter_message(message, reason="PermanentProcessingError",
                                             error_description=str(exc)[:1000])
            except Exception as exc:  # transient or unknown: Service Bus dead-letters it after the max delivery count
                logger.warning("Retrying message %s later: %s", message.message_id, exc)
                receiver.abandon_message(message)


if __name__ == "__main__":
    run_worker()