"""One short prompt, greedy, few tokens — a fast semantic sanity check on SM90.

`offline_inference.py` runs five prompts for up to 128 tokens each, which is minutes of decode
per attempt. This asks a single question and stops after `--max_tokens`, so a bring-up run tells
us what we actually care about: whether the first tokens are words rather than noise.

    python scripts/h20_ask.py "问题：中国的首都是哪里？\n答案：" --max_tokens 8
"""

import argparse
import os

from _config import DEFAULT_MODEL_NAME, create_llm, setup_common_env
from _launch import launch_dp

from deepinfer.sampling_params import SamplingParams

parser = argparse.ArgumentParser()
parser.add_argument(
    "prompt",
    nargs="*",
    default=["问题：中国的首都是哪里？\n答案：", "问题：中国的首都是哪里？\n答案：北京"],
)
parser.add_argument("--max_tokens", type=int, default=8)
parser.add_argument("--ep", type=int, default=8)
parser.add_argument("--cp", type=int, default=1)
parser.add_argument("--model_name", default=DEFAULT_MODEL_NAME)
parser.add_argument("--temperature", type=float, default=0.0, help="0 = greedy, so the answer is reproducible")
parser.add_argument(
    "--ids",
    default="",
    help="Comma-separated prompt token ids. Use this to feed *exactly* the ids the reference model saw "
    "(the official generate.py wraps text in a chat template, so raw strings are not comparable).",
)
args = parser.parse_args()

ep_size, cp_size = args.ep, args.cp
dp_size = max(ep_size // cp_size, 1)
world_size = dp_size * cp_size
assert world_size % ep_size == 0


def main(dp_size: int, dp_rank: int) -> None:
    setup_common_env(
        model_name=args.model_name,
        sampler_force_topk=4096,
        set_override_experts=False,
        enable_eplb=False,
    )
    os.environ.setdefault("DEEPINFER_BATCH_INVARIANT", "1")

    llm = create_llm(
        args.model_name,
        dp_rank,
        ep_size=ep_size,
        cp_size=cp_size,
        env={"DEEPINFER_LOGGING_INTERVAL_STEP": "10"},
        data_parallel_size=dp_size,
        max_num_batched_tokens=4096,
        max_model_len=4096,
        max_num_seqs=max(1, len(args.prompt)),
        gpu_memory_utilization=0.85,
        lookahead_tokens=0,
    )

    sampling_params = [
        SamplingParams(
            temperature=args.temperature,
            max_tokens=args.max_tokens,
            ignore_eos=True,
            top_logprobs=4,
            seq_seed=0,
        )
        for _ in args.prompt
    ]

    if args.ids:
        ids = [int(x) for x in args.ids.replace(" ", "").split(",") if x]
        outputs = llm.generate(
            prompts=None, prompt_token_ids=[ids], sampling_params=sampling_params[:1], use_tqdm=False
        )
        shown = [f"<ids len={len(ids)}>"]
    else:
        outputs = llm.generate(prompts=args.prompt, sampling_params=sampling_params, use_tqdm=False)
        shown = args.prompt
    for i, output in enumerate(outputs):
        assert output, f"sample {i}: no generated tokens"
        tokens = [x.token_id for x in output]
        logps = [x.token_logprobs for x in output]
        text = llm.llm_engine.tokenizer.decode(tokens)
        label = shown[i] if args.ids else args.prompt[i]
        print(f"\n=== DP {dp_rank} sample {i}")
        print(f"=== prompt: {label!r}")
        print(f"=== token ids ({len(tokens)}): {tokens}")
        print(f"=== per-token logprob: {[round(x, 3) for x in logps]}")
        print(f"=== continuation: {text!r}", flush=True)
        if args.ids:
            print(f"=== ids in / out: {len(ids)} -> {tokens}")

    llm.close()


if __name__ == "__main__":
    launch_dp(main, dp_size)