import io
import wave

import numpy as np
import polars as pl

from vllm import LLM, SamplingParams


def wav_bytes_to_numpy(wav_bytes):
    with wave.open(io.BytesIO(wav_bytes), "rb") as wr:
        if (num_channels := wr.getnchannels()) != 1:
            raise ValueError(f"num_channels must be 1, got {num_channels}")
        if (sample_width := wr.getsampwidth()) != 2:
            raise ValueError(f"sample_width must be 2, got {sample_width}")

        audio_data = wr.readframes(wr.getnframes())

        return np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0


llm = LLM(
    model="Yehor/whisper-large-v2-quantized-uk",
    max_model_len=448,
    max_num_seqs=400,
    gpu_memory_utilization=0.8,
    limit_mm_per_prompt={"audio": 1},
)

df = pl.read_parquet("hf://datasets/Yehor/cv10-uk-testset-clean/data/train-*.parquet")


for row in df.iter_rows(named=True):
    current_sample = (
        wav_bytes_to_numpy(row["audio"]["bytes"]),
        16_000,
    )

    inputs = {
        "encoder_prompt": {
            "prompt": "",
            "multi_modal_data": {
                "audio": current_sample,
            },
        },
        "decoder_prompt": "<|startoftranscript|><|uk|><|transcribe|><|notimestamps|>",
    }

    sampling_params = SamplingParams(
        temperature=0,
        top_p=1.0,
        max_tokens=200,
    )
    outputs = llm.generate(inputs, sampling_params)

    print(f"PROMPT  : {outputs[0].prompt}")
    print(f"TRANSCRIPTION: {row['transcription']}")
    print(f"PREDICTION: {outputs[0].outputs[0].text}")
    print("==========================================")