from dataclasses import dataclass, field
import logging

import gradio as gr
import torch
import transformers
import torchaudio

from multi_token.model_utils import MultiTaskType
from multi_token.training import ModelArguments
from multi_token.inference import load_trained_lora_model
from multi_token.data_tools import encode_chat


@dataclass
class ServeArguments(ModelArguments):
    load_bits: int = field(default=16)
    max_new_tokens: int = field(default=128)
    temperature: float = field(default=0.01)


# Load arguments and model
logging.getLogger().setLevel(logging.INFO)

parser = transformers.HfArgumentParser((ServeArguments,))
serve_args, _ = parser.parse_args_into_dataclasses(return_remaining_strings=True)

model, tokenizer = load_trained_lora_model(
    model_name_or_path=serve_args.model_name_or_path,
    model_lora_path=serve_args.model_lora_path,
    load_bits=serve_args.load_bits,
    use_multi_task=MultiTaskType(serve_args.use_multi_task),
    tasks_config=serve_args.tasks_config
)


def generate_caption(audio_file):
    waveform, sample_rate = torchaudio.load(audio_file)

    req_json = {
        "audio": {
            "tensor": waveform,
            "sampling_rate": sample_rate,
        }
    }

    encoded_dict = encode_chat(req_json, tokenizer, model.modalities)

    with torch.inference_mode():
        output_ids = model.generate(
            input_ids=encoded_dict["input_ids"].unsqueeze(0).to(model.device),
            max_new_tokens=serve_args.max_new_tokens,
            use_cache=True,
            do_sample=True,
            temperature=serve_args.temperature,
            modality_inputs={
                m.name: [encoded_dict[m.name]] for m in model.modalities
            },
        )

    outputs = tokenizer.decode(
        output_ids[0, encoded_dict["input_ids"].shape[0]:],
        skip_special_tokens=True
    ).strip()

    return outputs


demo = gr.Interface(
    fn=generate_caption,
    inputs=gr.Audio(type="filepath", label="Upload a WAV file"),
    outputs=gr.Textbox(label="Generated Caption"),
    title="Audio Caption Generator",
    description="Upload a .wav audio file to generate a caption using a LoRA fine-tuned model."
)

if __name__ == "__main__":
    demo.launch()