#!/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.

from __future__ import annotations

import argparse
import json
import os
import shutil
import sys
from pathlib import Path
from typing import Optional

import torch
from transformers import AutoConfig, AutoFeatureExtractor, AutoModelForCausalLM, AutoTokenizer

SCRIPT_DIR = Path(__file__).resolve().parent
if str(SCRIPT_DIR) not in sys.path:
    sys.path.insert(0, str(SCRIPT_DIR))

from audio_utils import (
    IM_END_TOKEN,
    build_attention_mask,
    build_prompt_template,
    expand_sound_placeholder,
    extract_whisper_features,
    load_audio,
    parse_conversation,
    resolve_audio_preprocessor_path,
    save_results_jsonl,
    split_thinking,
)


def refresh_remote_code_cache(model_path: str) -> None:
    """Drop stale HF dynamic-module cache for this local checkpoint folder."""
    module_name = Path(model_path).resolve().name.replace("-", "_").replace(".", "_")
    cache_root = os.environ.get(
        "HF_MODULES_CACHE",
        os.path.expanduser("~/.cache/huggingface/modules"),
    )
    cache_path = Path(cache_root) / "transformers_modules" / module_name
    if cache_path.exists():
        shutil.rmtree(cache_path)


def resolve_device_map(device_map: str, device: str):
    if device_map == "single":
        return {"": device}
    if device_map == "none":
        return None
    return device_map


def load_model(
    model_path: str,
    device_map: str,
    device: str,
    torch_dtype: str,
    tp_plan: str,
    refresh_code_cache: bool,
):
    print("loading:", model_path)
    if refresh_code_cache:
        refresh_remote_code_cache(model_path)
    dtype = getattr(torch, torch_dtype) if torch_dtype != "auto" else "auto"
    tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)
    print("tokenizer loaded")
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    print("config loaded")
    feature_extractor = AutoFeatureExtractor.from_pretrained(resolve_audio_preprocessor_path(model_path, config))
    print("feature_extractor loaded")
    tp_plan_arg = None if tp_plan == "none" else tp_plan
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        trust_remote_code=True,
        torch_dtype=dtype,
        device_map=resolve_device_map(device_map, device),
        tp_plan=tp_plan_arg,
    )
    print("model loaded")
    model.eval()
    return model, tokenizer, feature_extractor, config


def build_inputs(
    tokenizer,
    feature_extractor,
    config,
    audio_path: str,
    prompt: str,
    sample_rate: int,
    reasoning: bool,
    prompt_repitition: str,
):
    audio, sr = load_audio(audio_path, target_sr=sample_rate)
    input_features = extract_whisper_features(
        feature_extractor,
        audio,
        sample_rate=sr,
        clip_duration=float(getattr(config, "sound_clip_duration", 30.0)),
    )
    num_embeddings = input_features.shape[0] * int(getattr(config, "sound_embedding_size", 750))
    formatted_prompt = build_prompt_template(
        prompt,
        reasoning=reasoning,
        prompt_repitition=prompt_repitition,
    )
    expanded_prompt = expand_sound_placeholder(formatted_prompt, num_embeddings)
    tokenized = tokenizer(expanded_prompt, return_tensors="pt", add_special_tokens=False)
    input_ids = tokenized.input_ids
    attention_mask = tokenized.attention_mask if "attention_mask" in tokenized else build_attention_mask(input_ids)
    return input_ids, attention_mask, input_features, expanded_prompt


def generate_one(
    model,
    tokenizer,
    feature_extractor,
    config,
    audio_path: str,
    prompt: str,
    sample_rate: int,
    reasoning: bool,
    prompt_repitition: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
) -> str:
    input_ids, attention_mask, input_features, _ = build_inputs(
        tokenizer=tokenizer,
        feature_extractor=feature_extractor,
        config=config,
        audio_path=audio_path,
        prompt=prompt,
        sample_rate=sample_rate,
        reasoning=reasoning,
        prompt_repitition=prompt_repitition,
    )

    device = model.device
    input_ids = input_ids.to(device)
    attention_mask = attention_mask.to(device)
    input_features = input_features.to(device)

    eos_token_id = tokenizer.convert_tokens_to_ids(IM_END_TOKEN)
    if eos_token_id is None or eos_token_id == tokenizer.unk_token_id:
        eos_token_id = getattr(config, "eos_token_id", None)

    if temperature <= 0:
        raise ValueError(f"temperature must be > 0, got {temperature}")
    if not 0.0 <= top_p <= 1.0:
        raise ValueError(f"top_p must be in [0, 1], got {top_p}")
    if top_k < 0:
        raise ValueError(f"top_k must be >= 0, got {top_k}")

    do_sample = temperature != 1.0 or 0.0 < top_p < 1.0 or top_k > 0
    generation_kwargs = {
        "do_sample": do_sample,
        "eos_token_id": eos_token_id,
        "pad_token_id": tokenizer.pad_token_id or getattr(config, "pad_token_id", 0),
    }
    if do_sample:
        generation_kwargs["temperature"] = temperature
        if top_p > 0.0:
            generation_kwargs["top_p"] = top_p
        if top_k > 0:
            generation_kwargs["top_k"] = top_k

    with torch.inference_mode():
        output_ids = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            input_features=input_features,
            max_new_tokens=max_new_tokens,
            **generation_kwargs,
        )

    new_tokens = output_ids[0, input_ids.shape[-1] :]
    response = tokenizer.decode(new_tokens, skip_special_tokens=False)
    return response.split(IM_END_TOKEN, 1)[0].strip()


def run_inference(
    model,
    tokenizer,
    feature_extractor,
    config,
    input_samples: list[dict],
    audio_base_dir: Optional[str],
    sample_rate: int,
    reasoning: bool,
    prompt_repitition: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    top_k: int,
) -> list[dict]:
    results = []
    for idx, sample in enumerate(input_samples):
        sample_id = sample.get("id", idx)
        audio_path = sample["sound"]
        if audio_base_dir and not os.path.isabs(audio_path):
            audio_path = os.path.join(audio_base_dir, audio_path)

        human_prompt, gt_answer = parse_conversation(sample["conversations"])
        print(f"[{idx + 1}/{len(input_samples)}] id={sample_id} audio={audio_path}")

        response = generate_one(
            model=model,
            tokenizer=tokenizer,
            feature_extractor=feature_extractor,
            config=config,
            audio_path=audio_path,
            prompt=human_prompt,
            sample_rate=sample_rate,
            reasoning=reasoning,
            prompt_repitition=prompt_repitition,
            max_new_tokens=max_new_tokens,
            temperature=temperature,
            top_p=top_p,
            top_k=top_k,
        )

        thinking, prediction = split_thinking(response)
        result = {
            "id": sample_id,
            "sound": sample["sound"],
            "prompt": human_prompt,
            "ground_truth": gt_answer,
            "thinking": thinking,
            "prediction": prediction,
        }
        for key, value in sample.items():
            if key not in {"id", "sound", "conversations"}:
                result[key] = value
        results.append(result)
    return results


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--hf-model-path", required=True)
    parser.add_argument("--input-json", required=True)
    parser.add_argument("--output-jsonl", required=True)
    parser.add_argument("--audio-base-dir", default=None)
    parser.add_argument("--sample-rate", type=int, default=16000)
    parser.add_argument("--reasoning", action="store_true")
    parser.add_argument("--prompt-repitition", default="none", choices=["none", "repetition"], help="Prompt repitition trick.")
    parser.add_argument("--start-idx", type=int, default=0)
    parser.add_argument("--end-idx", type=int, default=None)
    parser.add_argument("--max-new-tokens", type=int, default=1024)
    parser.add_argument("--temperature", type=float, default=1.0)
    parser.add_argument(
        "--top-p",
        type=float,
        default=1.0,
        help="Top-p sampling threshold. Use 0.0 or 1.0 to disable top-p filtering.",
    )
    parser.add_argument(
        "--top-k",
        type=int,
        default=0,
        help="Top-k sampling threshold. Use 0 to disable top-k filtering.",
    )
    parser.add_argument(
        "--device-map",
        default="single",
        help="Default 'single' loads the whole model on --device. Use 'auto' for HF sharding or 'none' for CPU/default loading.",
    )
    parser.add_argument("--device", default="cuda:0", help="Device used when --device-map single.")
    parser.add_argument(
        "--tp-plan",
        default="none",
        choices=["none", "auto"],
        help="Use 'auto' only when launching with torch.distributed initialized for all TP ranks.",
    )
    parser.add_argument("--torch-dtype", default="bfloat16", choices=["auto", "float16", "bfloat16", "float32"])
    parser.add_argument(
        "--no-refresh-remote-code-cache",
        action="store_true",
        help="Do not clear this checkpoint's stale Hugging Face dynamic-module cache before loading.",
    )
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    model, tokenizer, feature_extractor, config = load_model(
        args.hf_model_path,
        device_map=args.device_map,
        device=args.device,
        torch_dtype=args.torch_dtype,
        tp_plan=args.tp_plan,
        refresh_code_cache=not args.no_refresh_remote_code_cache,
    )

    with open(args.input_json, "r", encoding="utf-8") as f:
        samples = json.load(f)
    end_idx = args.end_idx if args.end_idx is not None else len(samples)
    samples = samples[args.start_idx : end_idx]

    results = run_inference(
        model=model,
        tokenizer=tokenizer,
        feature_extractor=feature_extractor,
        config=config,
        input_samples=samples,
        audio_base_dir=args.audio_base_dir,
        sample_rate=args.sample_rate,
        reasoning=args.reasoning,
        prompt_repitition=args.prompt_repitition,
        max_new_tokens=args.max_new_tokens,
        temperature=args.temperature,
        top_p=args.top_p,
        top_k=args.top_k,
    )
    save_results_jsonl(results, args.output_jsonl)
    print(f"Saved {len(results)} results to {args.output_jsonl}")


if __name__ == "__main__":
    main()

    """
    CUDA_VISIBLE_DEVICES=0 python inference_hf.py \
        --hf-model-path checkpoint_folder_full/ \
        --input-json test_input.json \
        --output-jsonl outputs/output.jsonl \
        --start-idx 0 \
        --end-idx 16 \
        --max-new-tokens 1024 \
        --temperature 0.7 \
        --top-p 0.9 \
        --top-k 0
    """