import gradio as gr
import json
from nameder import init_model_ner, get_entity_results
from speech2text import init_model_trans, transcribe
from translation import translate
from resources import NER_Request, entity_labels_sample, set_start, audit_elapsedtime
import ast
import numpy as np

def translation_to_english(text: str):
   resultado = translate(text)
   return resultado

def transcription(audio):
    
    s2t = init_model_trans()
    sr, y = audio
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))
    return transcribe({"sampling_rate": sr, "raw": y}, s2t)

def named_entity_recognition(req: NER_Request):
    ner = init_model_ner()
    result = get_entity_results(entities_list=entity_labels_sample, 
                                model=ner,
                                text=req.text)
    print('result:',result,type(result))
    return json.dumps(result)

def get_lead(audio: bytes, labels: str, input_text: str):
    print("audio",audio,type(audio))
    print("input text:",input_text)
    print("labels:",labels)
    start = set_start()
    labels_list = ast.literal_eval(labels)
    if audio == None:
        text = input_text
    else:
        transcribe = transcription(audio)
        text = translation_to_english(transcribe)
        lead_input.value = text    
    ner = named_entity_recognition(NER_Request(
        entities=labels_list,
        text=text
    ))
    audit_elapsedtime("VoiceLead", start)
    return ner

audio_input = gr.Audio(
    label="Record your audio"
)
labels_input = gr.Textbox(
    label="Labels",
    info="Choose your labels",
    value=entity_labels_sample
)
lead_input = gr.Textbox(
            label="Lead",
            info="[Optional] Input your lead",
            lines=9,
            value="I have a lead that Salesforce needs 3 developers for 600 euros a day, for 6 months"
        )
text_output = gr.Textbox(
            label="Labels",
            info="",
            lines=9,
            value=""
        )
ui = gr.Interface(
    fn=get_lead,
    description= "Voice your lead",
    inputs=[audio_input, labels_input, lead_input],
    outputs=[text_output],
    title="VoiceLead"
)

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