"""Delete every chunk of a withdrawn source file from an Azure AI Search index.

Three ways to run it:
  --dry-run <source_path>  List the key of every chunk of the file, one per line. Deletes nothing.
  --keys-file <file>       Delete exactly the keys listed in the file (for example the dry-run output).
  <source_path>            Find the chunks of the file and delete them in one go.

Run --dry-run first, check the list, then delete with --keys-file, so you delete only what you checked.
Batches hold at most 1,000 documents (Azure AI Search service limits). Set BATCH_SIZE for smaller batches.

Prerequisites:
  pip install azure-search-documents azure-identity

Environment variables:
  SEARCH_ENDPOINT: Full search service endpoint (e.g. https://${SEARCH_SERVICE}.search.windows.net)
  INDEX_NAME: Target search index name
  KEY_FIELD: Document key field in the index (default: "id")
  PARENT_FIELD: Field storing the source document path or URI (default: "source_path")
  BATCH_SIZE: Documents per delete batch, 1 to 1000 (default: 1000)
"""
import argparse
import os
import sys
from typing import List

MAX_BATCH_SIZE = 1000  # Azure AI Search limit: maximum 1,000 documents per batch


def get_client():
    endpoint = os.environ.get("SEARCH_ENDPOINT")
    index_name = os.environ.get("INDEX_NAME")

    if not endpoint or not index_name:
        print("Error: SEARCH_ENDPOINT and INDEX_NAME environment variables must be set.", file=sys.stderr)
        sys.exit(1)

    try:
        from azure.identity import DefaultAzureCredential
        from azure.search.documents import SearchClient
    except ImportError as exc:
        print(f"Error: Required Azure SDK packages are missing ({exc}). Run: pip install azure-search-documents azure-identity", file=sys.stderr)
        sys.exit(1)

    return SearchClient(endpoint=endpoint, index_name=index_name, credential=DefaultAzureCredential())


def find_keys(client, source_path: str, key_field: str, parent_field: str) -> List[str]:
    # Escape single quotes for OData literal comparison
    safe_source = source_path.replace("'", "''")
    filter_expr = f"{parent_field} eq '{safe_source}'"

    print(f"Searching for chunks where {filter_expr}...", file=sys.stderr)
    results = client.search(search_text="*", filter=filter_expr, select=[key_field])
    return [doc[key_field] for doc in results if key_field in doc]


def read_keys(path: str) -> List[str]:
    # strip() also removes Windows line endings, so a key file written on Windows still works
    with open(path, encoding="utf-8") as handle:
        return [line.strip() for line in handle if line.strip()]


def delete_keys(client, keys: List[str], key_field: str, batch_size: int) -> List[str]:
    failed_keys: List[str] = []
    for start in range(0, len(keys), batch_size):
        documents_to_delete = [{key_field: k} for k in keys[start:start + batch_size]]
        for result in client.delete_documents(documents=documents_to_delete):
            if not result.succeeded:
                print(f"Failed to delete key: {result.key} (status: {result.status_code}, error: {result.error_message})", file=sys.stderr)
                failed_keys.append(result.key)
    return failed_keys


def main() -> None:
    parser = argparse.ArgumentParser(description="Delete every chunk of a withdrawn source file.")
    mode = parser.add_mutually_exclusive_group()
    mode.add_argument("--dry-run", action="store_true", help="list the chunk keys and delete nothing")
    mode.add_argument("--keys-file", help="delete exactly the keys in this file, one per line")
    parser.add_argument("source_path", nargs="?", help="value of the parent field, for example the blob URL")
    args = parser.parse_args()

    if not args.keys_file and not args.source_path:
        parser.error("give a source_path, or --keys-file")

    key_field = os.environ.get("KEY_FIELD", "id")
    parent_field = os.environ.get("PARENT_FIELD", "source_path")
    try:
        batch_size = int(os.environ.get("BATCH_SIZE", str(MAX_BATCH_SIZE)))
    except ValueError:
        batch_size = 0
    if not 1 <= batch_size <= MAX_BATCH_SIZE:
        print(f"Error: BATCH_SIZE must be a whole number from 1 to {MAX_BATCH_SIZE}.", file=sys.stderr)
        sys.exit(1)

    client = get_client()
    if args.keys_file:
        keys = read_keys(args.keys_file)
    else:
        keys = find_keys(client, args.source_path, key_field, parent_field)

    if args.dry_run:
        # Keys go to stdout and the summary to stderr, so "> chunks_to_delete.txt" holds only keys
        for key in keys:
            print(key)
        print(f"Found {len(keys)} chunks. Nothing deleted (dry run).", file=sys.stderr)
        return

    if not keys:
        print("No chunks to delete.")
        return

    print(f"Deleting {len(keys)} chunks in batches of {batch_size}...", file=sys.stderr)
    failed_keys = delete_keys(client, keys, key_field, batch_size)
    print(f"Successfully deleted {len(keys) - len(failed_keys)} of {len(keys)} chunks.")

    if failed_keys:
        print(f"Error: {len(failed_keys)} chunk deletions failed. Retry later for these keys: {failed_keys}", file=sys.stderr)
        sys.exit(1)


if __name__ == "__main__":
    main()