# 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 json
import math
import os
from pathlib import Path
from typing import Iterable, Optional

import numpy as np
import torch


SOUND_PLACEHOLDER = "<sound>"
SOUND_TOKEN = "<so_embedding>"
SOUND_START_TOKEN = "<so_start>"
SOUND_END_TOKEN = "<so_end>"
IM_END_TOKEN = "<|im_end|>"
DEFAULT_SYSTEM_PROMPT = (
    "<|im_start|>system\n"
    "You are a helpful and harmless assistant.\n\n"
    "You are not allowed to use any tools."
    "<|im_end|>\n"
)


def strip_hf_prefix(path: str) -> str:
    """Convert Megatron-style hf:// paths into local filesystem paths."""
    return path[len("hf://") :] if path.startswith("hf://") else path


def load_audio(audio_path: str, target_sr: int = 16000) -> tuple[np.ndarray, int]:
    import librosa

    audio_data, sr = librosa.load(audio_path, sr=target_sr, mono=True)
    return normalize_audio(audio_data), sr


def normalize_audio(audio: np.ndarray) -> np.ndarray:
    """Return mono float32 audio in [-1, 1], matching the Megatron eval path."""
    audio = np.asarray(audio)
    if audio.ndim == 2:
        if audio.shape[1] <= 2:
            audio = audio.mean(axis=1)
        elif audio.shape[0] <= 2:
            audio = audio.mean(axis=0)
        else:
            raise ValueError(f"Unsupported audio shape: {audio.shape}")

    if audio.dtype == np.int16:
        audio = audio.astype(np.float32) / 32768.0
    elif audio.dtype != np.float32:
        audio = audio.astype(np.float32)

    max_abs = float(np.abs(audio).max()) if audio.size else 0.0
    if max_abs > 1.0:
        audio = audio / max_abs
    return audio.astype(np.float32, copy=False)


def split_audio_into_clips(
    audio: np.ndarray,
    sample_rate: int = 16000,
    clip_duration: float = 30.0,
) -> list[np.ndarray]:
    """Split audio into fixed 30s clips; keep a padded final clip for Whisper."""
    audio = normalize_audio(audio)
    clip_samples = int(round(sample_rate * clip_duration))
    if clip_samples <= 0:
        raise ValueError(f"Invalid clip_samples: {clip_samples}")
    if audio.size == 0:
        audio = np.zeros(1, dtype=np.float32)

    num_clips = max(1, math.ceil(audio.shape[0] / clip_samples))
    clips: list[np.ndarray] = []
    for idx in range(num_clips):
        start = idx * clip_samples
        clip = audio[start : start + clip_samples]
        if clip.shape[0] < clip_samples:
            clip = np.pad(clip, (0, clip_samples - clip.shape[0]))
        clips.append(clip.astype(np.float32, copy=False))
    return clips


def extract_whisper_features(
    feature_extractor,
    audio: np.ndarray,
    sample_rate: int = 16000,
    clip_duration: float = 30.0,
) -> torch.Tensor:
    """Return NV-Whisper input features shaped (num_clips, 128, 3000)."""
    clips = split_audio_into_clips(audio, sample_rate=sample_rate, clip_duration=clip_duration)
    features = feature_extractor(
        clips,
        sampling_rate=sample_rate,
        return_tensors="pt",
        padding="max_length",
        return_attention_mask=False,
    )
    input_features = features.input_features
    if input_features.ndim != 3:
        raise ValueError(f"Expected 3D Whisper features, got {tuple(input_features.shape)}")
    return input_features


def parse_conversation(conversation: list[dict]) -> tuple[str, str]:
    human_prompt = ""
    gt_answer = ""
    for turn in conversation:
        if turn["from"] == "human":
            human_prompt = turn["value"].replace("<sound>\n", "").replace("<sound>", "").strip()
        elif turn["from"] == "gpt":
            gt_answer = turn["value"]
    return human_prompt, gt_answer


def build_prompt_template(
    prompt: str,
    reasoning: bool = False,
    prompt_repitition: str = "none",
) -> str:
    if prompt_repitition not in {"none", "repetition"}:
        raise ValueError(f"Unknown prompt repetition mode: {prompt_repitition}")
    if prompt_repitition == "repetition":
        prompt = f"{prompt}\n{prompt}"

    if reasoning:
        return f"<|im_start|>user\n<sound>\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think>\n"
    return f"<|im_start|>user\n<sound>\n{prompt}<|im_end|>\n<|im_start|>assistant\n<think></think>"


def expand_sound_placeholder(prompt: str, num_embeddings: int) -> str:
    if prompt.count(SOUND_PLACEHOLDER) != 1:
        raise ValueError(f"Expected exactly one {SOUND_PLACEHOLDER}, found {prompt.count(SOUND_PLACEHOLDER)}")
    replacement = SOUND_START_TOKEN + (SOUND_TOKEN * num_embeddings) + SOUND_END_TOKEN
    return prompt.replace(SOUND_PLACEHOLDER, replacement)


def build_attention_mask(input_ids: torch.Tensor) -> torch.Tensor:
    return torch.ones_like(input_ids, dtype=torch.long)


def split_thinking(response: str) -> tuple[str, str]:
    if "</think>" not in response:
        return "", response.strip()
    thinking = response.rsplit("</think>", 1)[0].strip() + "</think>"
    prediction = response.rsplit("</think>", 1)[1].strip()
    return thinking, prediction


def save_results_jsonl(results: Iterable[dict], output_path: str) -> None:
    os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
    with open(output_path, "w", encoding="utf-8") as f:
        for result in results:
            f.write(json.dumps(result, ensure_ascii=False) + "\n")


def resolve_audio_preprocessor_path(model_path: str, config) -> str:
    path = getattr(config, "audio_preprocessor_path", None) or "audio_preprocessor"
    candidate = Path(path)
    if not candidate.is_absolute():
        candidate = Path(model_path) / candidate
    return str(candidate)