import gradio as gr
import os
import requests
import json

API_TOKEN = os.getenv('HF_API_TOKEN')

PROVIDERMODEL = {
    'openai': 'openai/whisper-medium.en',
    'facebook': 'facebook/wav2vec2-base-960h',
}

SUMMARYMODEL = {
    'openai': 'openai/whisper-medium.en',
    'facebook': 'facebook/bart-large-cnn',
}

def query(filename, provider='openai'):
    with open(filename, "rb") as f:
        data = f.read()
    API_URL = f"https://api-inference.huggingface.co/models/{PROVIDERMODEL[provider]}"
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    response = requests.request("POST", API_URL, headers=headers, data=data)
    if response.status_code == 200:
        return json.loads(response.content.decode("utf-8")).get('text')
    else:
        return response.content


def summarize(text, provider='facebook'):
    data = {
        'inputs': text,
        'parameters': {
            'do_sample': False
        }
    }
    API_URL = f"https://api-inference.huggingface.co/models/{SUMMARYMODEL[provider]}"
    headers = {"Authorization": f"Bearer {API_TOKEN}"}
    response = requests.request("POST", API_URL, headers=headers, json=data)
    if response.status_code == 200:
        return json.loads(response.content.decode("utf-8"))[0].get('summary_text')
    else:
        return {"error": response.content}


demo = gr.Blocks()

with demo:
    audio_file = gr.Audio(type="filepath")
    text = gr.Textbox()
    label = gr.Textbox()

    b1 = gr.Button("Derive Transcript")
    b2 = gr.Button("Summarize")

    b1.click(query, inputs=audio_file, outputs=text)
    b2.click(summarize, inputs=text, outputs=label)
demo.launch()