"""Text-only PitchIn-Q inference using the paper's merged checkpoint."""
import os
import time
from pathlib import Path

import spaces  # Must precede torch for ZeroGPU's CUDA initialization.
import torch
from peft import PeftModel
from transformers import AutoConfig, AutoProcessor, Qwen2_5OmniForConditionalGeneration

GENERATION = dict(
    do_sample=True, temperature=0.6, top_p=0.85, top_k=30,
    repetition_penalty=1.08, max_new_tokens=384,
    eos_token_id=151645, pad_token_id=151643,
)
_ENGINE = None


@spaces.GPU(duration=120)
def generate_on_gpu(wave, question):
    # Keep the startup model in the forked worker's globals. Passing a bound
    # method's `self` through ZeroGPU's argument queue reserializes it on CPU.
    if _ENGINE is None:
        raise RuntimeError("PitchIn is not initialized.")
    return _ENGINE._generate(wave, question)


class PitchIn:
    def __init__(self):
        global _ENGINE
        root = Path(__file__).parent
        merged_path = os.environ.get("PITCHIN_MODEL_PATH")
        model_path = merged_path or os.environ.get("PITCHIN_BASE_MODEL", "Qwen/Qwen2.5-Omni-7B")
        model_args = {"local_files_only": True} if Path(model_path).is_dir() else {
            "revision": "ae9e1690543ffd5c0221dc27f79834d0294cba00",
        }
        processor_path = merged_path or os.environ.get("PITCHIN_PROCESSOR_PATH", str(root / "processor"))
        config = AutoConfig.from_pretrained(model_path, **model_args)
        config.enable_audio_output = False
        self.processor = AutoProcessor.from_pretrained(processor_path, local_files_only=True)
        outer = Qwen2_5OmniForConditionalGeneration.from_pretrained(
            model_path, config=config, dtype=torch.bfloat16,
            attn_implementation="sdpa", **model_args,
        )
        if not merged_path:
            adapter_path = os.environ.get("PITCHIN_ADAPTER_PATH", str(root / "adapter"))
            outer.thinker = PeftModel.from_pretrained(
                outer.thinker, adapter_path, is_trainable=False, torch_device="cpu",
            ).merge_and_unload()
        self.model = outer.thinker.to("cuda").eval()
        if self.processor.feature_extractor.sampling_rate != 16000:
            raise RuntimeError("Unexpected checkpoint audio sample rate.")
        _ENGINE = self

    def generate(self, wave, question):
        return generate_on_gpu(wave, question)

    def _generate(self, wave, question):
        started = time.monotonic()
        messages = [{"role": "user", "content": [
            {"type": "audio", "audio": "uploaded_audio.wav"},
            {"type": "text", "text": question},
        ]}]
        text = self.processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        inputs = self.processor(
            text=[text], audio=[wave], sampling_rate=16000, return_tensors="pt",
            audio_kwargs={"truncation": False},
            text_kwargs={"padding": True, "truncation": False},
        )
        inputs = {k: v.to("cuda", dtype=torch.bfloat16 if v.is_floating_point() else v.dtype)
                  if isinstance(v, torch.Tensor) else v for k, v in inputs.items()}
        input_length = inputs["input_ids"].shape[1]
        with torch.inference_mode():
            tokens = self.model.generate(**inputs, **GENERATION)[0, input_length:]
        response = self.processor.tokenizer.decode(tokens, skip_special_tokens=True).strip()
        if not response:
            raise RuntimeError("The model returned an empty answer. Please try again.")
        count = len(tokens)
        stopped_at_limit = count == GENERATION["max_new_tokens"] and tokens[-1].item() != GENERATION["eos_token_id"]
        status = f"{len(wave) / 16000:.1f}s audio · {time.monotonic() - started:.1f}s generation"
        if stopped_at_limit:
            status += " · Response reached the demo length limit."
        return response, status