import gradio as gr
import os 
import json 
import requests
import time
 
# AssemblyAI transcript endpoint (where we submit the file)
transcript_endpoint = "https://api.assemblyai.com/v2/transcript"

upload_endpoint = "https://api.assemblyai.com/v2/upload"


# Helper function to upload data
def _read_file(filename, chunk_size=5242880):
    with open(filename, "rb") as f:
        while True:
            data = f.read(chunk_size)
            if not data:
                break
            yield data
            
def get_transcript_url(url, api_token):
    headers={
    "Authorization": api_token,
    "Content-Type": "application/json"
    }
    
    # JSON that tells the API which file to trancsribe
    json={
    # URL of the audio file to process
    "audio_url": url,

    # Turn on speaker labels
    "speaker_labels": True,

    # Turn on cusom vocabulary
    "word_boost": ["assembly ai"],

    # Turn on custom spelling
    "custom_spelling": [
        {"from": ["assembly AI"], "to": "AssemblyAI"},
        {"from": ["assembly AI's"], "to": "AssemblyAI's"}
        ],

    # Turn on PII Redaction and specify policies
    "redact_pii": True,
    "redact_pii_policies": ["drug", "injury", "person_name"],
    "redact_pii_audio": True,

    # Turn on Auto Highlights
    "auto_highlights": True,

    # Turn on Content Moderation
    "content_safety": True,

    # Turn on Topic Detection
    "iab_categories": True,

    # Turn on Sentiment Analysis
    "sentiment_analysis": True,

    # Turn on Summarization and specify configuration
    "summarization": True,
    "summary_model": "informative",
    "summary_type": "bullets",

    # Turn on Entity Detection
    "entity_detection": True,}

    response = requests.post(
        transcript_endpoint,
        json=json,
        headers=headers  # Authorization to link this transcription with your account
      )

    polling_endpoint = f"https://api.assemblyai.com/v2/transcript/{response.json()['id']}"
    while True:
      transcription_result = requests.get(polling_endpoint, headers=headers).json()
      if transcription_result['status'] == 'completed':
        break
      elif transcription_result['status'] == 'error':
        raise RuntimeError(f"Transcription failed: {transcription_result['error']}")
      else:
        time.sleep(3)
    return transcription_result['text']

def get_transcript_file(filename, api_token):
    headers={
    "Authorization": api_token,
    "Content-Type": "application/json"
    }
    
    upload_response = requests.post(
    upload_endpoint,
    headers=headers, 
    data=_read_file(filename))
    
    # JSON that tells the API which file to trancsribe
    json = {
    # URL of the audio file to process
    "audio_url": upload_response.json()['upload_url'],

    # Turn on speaker labels
    "speaker_labels": True,

    # Turn on cusom vocabulary
    "word_boost": ["assembly ai"],

    # Turn on custom spelling
    "custom_spelling": [
        {"from": ["assembly AI"], "to": "AssemblyAI"},
        {"from": ["assembly AI's"], "to": "AssemblyAI's"}
        ],

    # Turn on PII Redaction and specify policies
    "redact_pii": True,
    "redact_pii_policies": ["drug", "injury", "person_name"],
    "redact_pii_audio": True,

    # Turn on Auto Highlights
    "auto_highlights": True,

    # Turn on Content Moderation
    "content_safety": True,

    # Turn on Topic Detection
    "iab_categories": True,

    # Turn on Sentiment Analysis
    "sentiment_analysis": True,

    # Turn on Summarization and specify configuration
    "summarization": True,
    "summary_model": "informative",
    "summary_type": "bullets",

    # Turn on Entity Detection
    "entity_detection": True,
    }

    response = requests.post(
        transcript_endpoint,
        json=json,
        headers=headers  # Authorization to link this transcription with your account
      )

    polling_endpoint = f"https://api.assemblyai.com/v2/transcript/{response.json()['id']}"
    while True:
      transcription_result = requests.get(polling_endpoint, headers=headers).json()
      if transcription_result['status'] == 'completed':
        break
      elif transcription_result['status'] == 'error':
        raise RuntimeError(f"Transcription failed: {transcription_result['error']}")
      else:
        time.sleep(3)
    return transcription_result['text']


title = """<h1 align="center">🔥Conformer-1 API </h1>"""
description = """
## In this demo, you can explore the outputs of a Conformer-1 Speech Recognition Model from AssemblyAI.
"""
                
with gr.Blocks(css = """#col_container {width: 1000px; margin-left: auto; margin-right: auto;}
                """) as demo:
    gr.HTML(title)
    gr.Markdown(description)

    assemblyai_api_key = gr.Textbox(type='password', label="Enter your AssemblyAI API key here")

    with gr.Column(elem_id = "col_container"):
        
        with gr.Tab("Audio URL file"):
            inputs = gr.Textbox(label = "Enter the url for the audio file")        
            b1 = gr.Button('Transcribe')
        
        with gr.Tab("Upload Audio as File"):
            audio_input_u = gr.Audio(label = 'Upload Audio',source="upload",type="filepath")
            transcribe_audio_u = gr.Button('Transcribe')

        transcript = gr.Textbox(label = "Transcript Result" )
    
    inputs.submit(get_transcript_url, [inputs, assemblyai_api_key], [transcript])
    b1.click(get_transcript_url, [inputs, assemblyai_api_key], [transcript])
    audio_input_u.submit(get_transcript_file, [audio_input_u, assemblyai_api_key], [transcript])
    transcibe_audio_u.click(get_transcript_file, [audio_input_u, assemblyai_api_key], [transcript])
                    
    
    demo.queue().launch(debug=True)