"""Nightly reconciliation: compare the files in a blob container with the chunks in the index.

For the index that the Pattern 3 worker (worker.py) fills. Reports:
- orphans: chunks whose file no longer exists in the container
- missing files: files in the container that have no chunks in the index
By default it only reports. With --delete-orphans it also deletes the orphan chunks.
Exit code: 0 no drift, 1 a delete failed, 2 drift found.

Environment variables: SEARCH_ENDPOINT, INDEX_NAME, STORAGE_ACCOUNT, STORAGE_CONTAINER,
KEY_FIELD (default "id"; must be filterable and sortable), PARENT_FIELD (default "source_url")
Prerequisites: pip install azure-search-documents azure-storage-blob azure-identity
"""
import argparse
import os
import sys
from typing import Dict, List, Set
from urllib.parse import unquote

from azure.identity import DefaultAzureCredential
from azure.search.documents import SearchClient
from azure.storage.blob import ContainerClient

PAGE_SIZE = 1000  # Azure AI Search returns at most 1,000 documents per page, and takes 1,000 per batch


def index_chunks(search_client, key_field: str, parent_field: str) -> Dict[str, List[str]]:
    """Map each file URL to its chunk keys. Pages by key with orderby and a range filter."""
    by_parent: Dict[str, List[str]] = {}
    last_key = None
    while True:
        page_filter = None if last_key is None else f"{key_field} gt '{last_key.replace(chr(39), chr(39) * 2)}'"
        page = list(search_client.search(search_text="*", select=[key_field, parent_field],
                                         order_by=[f"{key_field} asc"], filter=page_filter, top=PAGE_SIZE))
        for doc in page:
            by_parent.setdefault(unquote(doc.get(parent_field) or ""), []).append(doc[key_field])
        if len(page) < PAGE_SIZE:
            return by_parent
        last_key = page[-1][key_field]


def container_urls(container_client) -> Set[str]:
    return {unquote(container_client.get_blob_client(blob.name).url) for blob in container_client.list_blobs()}


def main() -> None:
    parser = argparse.ArgumentParser(description="Compare a blob container with the chunks in the index.")
    parser.add_argument("--delete-orphans", action="store_true", help="also delete the chunks of files that no longer exist")
    args = parser.parse_args()

    key_field = os.environ.get("KEY_FIELD", "id")
    parent_field = os.environ.get("PARENT_FIELD", "source_url")
    credential = DefaultAzureCredential()
    search_client = SearchClient(os.environ["SEARCH_ENDPOINT"], os.environ["INDEX_NAME"], credential)
    account_url = f"https://{os.environ['STORAGE_ACCOUNT']}.blob.core.windows.net"
    container_client = ContainerClient(account_url, os.environ["STORAGE_CONTAINER"], credential=credential)

    chunks = index_chunks(search_client, key_field, parent_field)
    files = container_urls(container_client)
    orphans = {url: keys for url, keys in chunks.items() if url not in files}
    missing = sorted(files - set(chunks))

    for url, keys in sorted(orphans.items()):
        print(f"ORPHAN  {url} ({len(keys)} chunks)")
    for url in missing:
        print(f"MISSING {url}")
    orphan_keys = [key for keys in orphans.values() for key in keys]
    print(f"Orphan files: {len(orphans)}, orphan chunks: {len(orphan_keys)}, missing files: {len(missing)}")

    if args.delete_orphans and orphan_keys:
        failed = 0
        for start in range(0, len(orphan_keys), PAGE_SIZE):
            batch = [{key_field: key} for key in orphan_keys[start:start + PAGE_SIZE]]
            failed += sum(1 for r in search_client.delete_documents(documents=batch) if not r.succeeded)
        print(f"Deleted {len(orphan_keys) - failed} of {len(orphan_keys)} orphan chunks.")
        if failed:
            sys.exit(1)
        orphans = {}
    if orphans or missing:
        sys.exit(2)


if __name__ == "__main__":
    main()