#!/usr/bin/env python3
# coding=utf-8
# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""OpenAI-compatible client for the Audex-2B audio-understanding server.

Sends a local audio file as an `audio_url` (file:// URI) plus a text question,
mirroring the Nemotron-3 Nano Omni usage. The server must be started with
`serve_audioqa_vllm.sh`.

Benchmark reproduction is non-thinking by default (`--recipe
audio-understanding`): the request sets `chat_template_kwargs={"enable_thinking":
False}` and the recipe sampling values.

Generation is masked to text ids (`allowed_token_ids = range(131072)` minus the
sound placeholder ids), passed via `extra_body`, so audio codec/gen and `<so_*>`
placeholder tokens can never leak into the response.

Usage:
  python client_audioqa.py \
      --audio /path/to/audio.wav \
      --prompt "Briefly describe the main sound in this audio." \
      --base-url http://localhost:8000/v1
"""

from __future__ import annotations

import argparse
import json
from pathlib import Path

TEXT_VOCAB = 131072
LEAK_PATTERNS = ("<audiocodec_", "<speechcodec_", "<audiogen_", "<speechgen_",
                 "<so_embedding>", "<so_start>", "<so_end>")
DEFAULT_PLACEHOLDER_IDS = (29, 30, 31)

RECIPES = {
    "audio-understanding":            {"reasoning": False, "temperature": 0.7, "top_p": 0.9, "top_k": 0},
    "speech-recognition-translation": {"reasoning": False, "temperature": 0.0, "top_p": 1.0, "top_k": 0},
    "custom":                         {"reasoning": True,  "temperature": 0.7, "top_p": 0.9, "top_k": 0},
}


def placeholder_ids(model_path: str | None) -> tuple[int, ...]:
    if not model_path:
        return DEFAULT_PLACEHOLDER_IDS
    cfg = json.loads((Path(model_path) / "config.json").read_text())
    ids = [cfg.get("sound_token_id"), cfg.get("sound_start_token_id"), cfg.get("sound_end_token_id")]
    return tuple(int(i) for i in ids if i is not None) or DEFAULT_PLACEHOLDER_IDS


def resolve_recipe(args: argparse.Namespace) -> dict:
    cfg = dict(RECIPES[args.recipe])
    if args.reasoning is not None:
        cfg["reasoning"] = args.reasoning
    if args.temperature is not None:
        cfg["temperature"] = args.temperature
    if args.top_p is not None:
        cfg["top_p"] = args.top_p
    if args.top_k is not None:
        cfg["top_k"] = args.top_k
    return cfg


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--audio", required=True)
    p.add_argument("--prompt", default="Briefly describe the main sound in this audio.")
    p.add_argument("--base-url", default="http://localhost:8000/v1")
    p.add_argument("--model", default="audex_audio")
    p.add_argument("--model-path", default=None,
                   help="Optional checkpoint dir to read sound placeholder ids from; "
                        "defaults to the released ids (29/30/31).")
    p.add_argument("--max-tokens", type=int, default=1024)
    p.add_argument(
        "--recipe",
        choices=sorted(RECIPES),
        default="audio-understanding",
        help="Benchmark recipe (default: audio-understanding, non-thinking). "
             "Explicit flags below override recipe values.",
    )
    p.add_argument("--reasoning", action=argparse.BooleanOptionalAction, default=None,
                   help="Override recipe thinking mode (--reasoning / --no-reasoning); "
                        "maps to chat_template_kwargs.enable_thinking.")
    p.add_argument("--temperature", type=float, default=None, help="Override recipe temperature.")
    p.add_argument("--top-p", type=float, default=None, help="Override recipe top_p.")
    p.add_argument("--top-k", type=int, default=None, help="Override recipe top_k (0 = disabled).")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    from openai import OpenAI

    cfg = resolve_recipe(args)
    print(f"[recipe={args.recipe}] enable_thinking={cfg['reasoning']} "
          f"temperature={cfg['temperature']} top_p={cfg['top_p']} top_k={cfg['top_k']}")

    blocked = set(placeholder_ids(args.model_path))
    allowed_ids = [i for i in range(TEXT_VOCAB) if i not in blocked]

    client = OpenAI(base_url=args.base_url, api_key="EMPTY")
    audio_uri = Path(args.audio).resolve().as_uri()

    resp = client.chat.completions.create(
        model=args.model,
        messages=[
            {
                "role": "user",
                "content": [
                    {"type": "audio_url", "audio_url": {"url": audio_uri}},
                    {"type": "text", "text": args.prompt},
                ],
            }
        ],
        max_tokens=args.max_tokens,
        temperature=cfg["temperature"],
        top_p=cfg["top_p"],
        extra_body={
            "top_k": cfg["top_k"],
            "allowed_token_ids": allowed_ids,
            "chat_template_kwargs": {"enable_thinking": cfg["reasoning"]},
        },
    )
    content = resp.choices[0].message.content
    leaks = [p for p in LEAK_PATTERNS if p in (content or "")]
    if leaks:
        raise SystemExit(f"Audio/placeholder token leakage in response: {leaks}\n{content}")
    print(content)


if __name__ == "__main__":
    main()