"""Probe one question against an Azure AI Search index in three modes.

Runs keyword-only, vector-only and hybrid queries for one question and reports the rank of the
expected document in each. If you pass a filter, it also runs hybrid without the filter, so you
can see whether the filter hides the document. It then suggests a class for the case. The
suggestion is a guide: confirm it by hand (see PB-RAG-02-03).

Read-only: it only sends search queries. It signs in with Microsoft Entra ID and never uses an
admin key. The identity needs the Search Index Data Reader role on the index.

Usage:
  python probe_query.py "<question>" <expected-document-id> [options]

Environment variables:
  SEARCH_ENDPOINT     https://${SEARCH_SERVICE}.search.windows.net
  SEARCH_INDEX        the index or alias that the application queries
  VECTOR_FIELD        the vector field to search (default: content_vector)
  ID_FIELD            the field that holds the document id to compare (default: parent_id)
  EMBEDDING_ENDPOINT  only for the fallback below: https://${OPENAI_RESOURCE}.openai.azure.com

By default the vector query uses kind "text", so the service vectorizes the question. That needs
a vectorizer on the vector field.

Fallback for an index with no vectorizer: pass --embedding-deployment and
--embedding-api-version. The script then calls your Azure OpenAI embedding deployment itself,
with Microsoft Entra ID, and sends the vector. Use the SAME embedding model that built the index,
or the ranks are meaningless. The identity needs the Cognitive Services OpenAI User role on the
Azure OpenAI resource. Field names are placeholders: set them to match your index.
"""
import argparse
import json
import os
import sys

import requests
from azure.identity import DefaultAzureCredential

SEARCH_SCOPE = "https://search.azure.com/.default"
# Token audience for Azure OpenAI with Entra ID, as in the Azure OpenAI Entra ID how-to.
EMBEDDING_SCOPE = "https://ai.azure.com/.default"


def get_token(scope=SEARCH_SCOPE):
    return DefaultAzureCredential().get_token(scope).token


def embed(endpoint, deployment, api_version, token, text):
    """Return the embedding of the text from an Azure OpenAI embedding deployment."""
    url = f"{endpoint}/openai/deployments/{deployment}/embeddings?api-version={api_version}"
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    resp = requests.post(url, headers=headers, json={"input": text}, timeout=30)
    if not resp.ok:
        sys.exit(f"Embedding call failed: HTTP {resp.status_code}: {resp.text[:500]}")
    return resp.json()["data"][0]["embedding"]


def search(endpoint, index, api_version, token, body):
    url = f"{endpoint}/indexes/{index}/docs/search?api-version={api_version}"
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    resp = requests.post(url, headers=headers, json=body, timeout=30)
    if not resp.ok:
        sys.exit(f"Search failed: HTTP {resp.status_code}: {resp.text[:500]}")
    return resp.json()["value"]


def build_body(mode, args, vector_field, id_field, filtered=True, semantic=False, query_vector=None):
    """Build the request body for one mode: keyword, vector or hybrid."""
    body = {"select": id_field, "top": args.depth}
    if mode in ("keyword", "hybrid"):
        body["search"] = args.question
    if mode in ("vector", "hybrid"):
        if query_vector is None:  # the service vectorizes the text
            vector_query = {"kind": "text", "text": args.question}
        else:  # fallback: the script sends a vector that it embedded itself
            vector_query = {"kind": "vector", "vector": query_vector}
        body["vectorQueries"] = [{**vector_query, "fields": vector_field, "k": args.depth}]
    if args.filter and filtered:
        body["filter"] = args.filter
        if mode != "keyword" and args.filter_mode:
            body["vectorFilterMode"] = args.filter_mode
    if mode == "hybrid" and semantic:
        body["queryType"] = "semantic"
        body["semanticConfiguration"] = args.semantic
        if args.debug:
            body["debug"] = "all"
    elif mode == "hybrid" and args.debug and not args.semantic and filtered:
        body["debug"] = "all"
    return body


def rank_of(hits, id_field, expected):
    """Return the 1-based rank of the first hit with the expected id, and that hit."""
    for position, hit in enumerate(hits, start=1):
        if str(hit.get(id_field)) == expected:
            return position, hit
    return None, None


def classify(ranks, cutoff, has_filter):
    """Suggest a class from the ranks. This is a guide, not a verdict."""
    primary = "hybrid + semantic" if "hybrid + semantic" in ranks else "hybrid"
    hybrid = ranks[primary]
    if hybrid is not None and hybrid <= cutoff:
        return f"in the top {cutoff}: retrieval looks fine, check generation (step 7)"
    if hybrid is not None:
        return f"ranking issue: {primary} rank {hybrid} is below the top {cutoff}"
    others = [m for m in ("keyword", "vector") if ranks[m] is not None]
    if others:
        return f"ranking issue: {primary} misses it but {' and '.join(others)} found it"
    if has_filter and ranks.get("hybrid, no filter") is not None:
        return "filter issue: found only without the filter"
    return "retrieval miss: not found in any mode within the search depth"


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("question")
    parser.add_argument("expected_id", help="id of the document that should be returned")
    parser.add_argument("--filter", help="OData filter that the user's query applies")
    parser.add_argument(
        "--filter-mode",
        choices=["preFilter", "postFilter", "strictPostFilter"],
        help="vectorFilterMode; strictPostFilter is preview and needs a preview --api-version",
    )
    parser.add_argument(
        "--semantic",
        metavar="CONFIG",
        help="semantic configuration name; adds a hybrid + semantic run and judges the class on it",
    )
    parser.add_argument("--debug", action="store_true", help="print the debug hit for hybrid")
    parser.add_argument("--top", type=int, default=10, help="cutoff the application uses")
    parser.add_argument("--depth", type=int, default=50, help="how far down to look")
    parser.add_argument("--api-version", default="2026-04-01")
    parser.add_argument(
        "--embedding-deployment",
        help="fallback for an index with no vectorizer: embed the question with this Azure OpenAI "
        "deployment (the same model that built the index)",
    )
    parser.add_argument(
        "--embedding-api-version",
        help="Azure OpenAI api-version for the embedding call (required with --embedding-deployment)",
    )
    args = parser.parse_args()
    if args.embedding_deployment and not args.embedding_api_version:
        parser.error("--embedding-deployment needs --embedding-api-version")

    try:
        endpoint = os.environ["SEARCH_ENDPOINT"].rstrip("/")
        index = os.environ["SEARCH_INDEX"]
        embedding_endpoint = (
            os.environ["EMBEDDING_ENDPOINT"].rstrip("/") if args.embedding_deployment else None
        )
    except KeyError as missing:
        sys.exit(f"Set the environment variable {missing}")
    vector_field = os.environ.get("VECTOR_FIELD", "content_vector")
    id_field = os.environ.get("ID_FIELD", "parent_id")
    token = get_token()
    query_vector = None
    if args.embedding_deployment:
        query_vector = embed(
            embedding_endpoint,
            args.embedding_deployment,
            args.embedding_api_version,
            get_token(EMBEDDING_SCOPE),
            args.question,
        )

    runs = [("keyword", "keyword", {}), ("vector", "vector", {}), ("hybrid", "hybrid", {})]
    if args.semantic:
        runs.append(("hybrid + semantic", "hybrid", {"semantic": True}))
    if args.filter:
        runs.append(("hybrid, no filter", "hybrid", {"filtered": False}))

    ranks, debug_hit = {}, None
    for label, mode, options in runs:
        body = build_body(mode, args, vector_field, id_field, query_vector=query_vector, **options)
        hits = search(endpoint, index, args.api_version, token, body)
        ranks[label], hit = rank_of(hits, id_field, args.expected_id)
        if hit is not None and body.get("debug"):
            debug_hit = hit

    print(f"Question:    {args.question}")
    print(f"Expected id: {args.expected_id}")
    print(f"{'Mode':<20}Rank (within the first {args.depth})")
    for label, _mode, _options in runs:
        print(f"{label:<20}{ranks[label] if ranks[label] is not None else '-'}")
    print("Suggested class:", classify(ranks, args.top, bool(args.filter)))
    if debug_hit is not None:
        print("Debug hit:", json.dumps(debug_hit, indent=2))


if __name__ == "__main__":
    main()