import gradio as gr
from transformers import pipeline
import numpy as np
import os
from huggingface_hub import login
from scipy.io.wavfile import write
import uuid
import torch
import boto3

print("cuda", torch.cuda.is_available())

access_token_read = os.environ.get('HF_TOKEN', None)
access_key = os.environ.get('access_key', None)
secret_access_key = os.environ.get('secret_access_key', None)

login(token = access_token_read)

session = boto3.Session(
    aws_access_key_id=access_key,
    aws_secret_access_key=secret_access_key,
)

s3 = session.resource('s3')

BUCKET = "audio-text-938"

print("cur path", os.listdir(os.path.join("..", "..", "..")))

if not os.path.isdir(os.path.join("..", "..", "..", "data", "hfcache")):
    os.mkdir(os.path.join("..", "..", "..", "data", "hfcache"))
    
if not os.path.isdir(os.path.join("..", "..", "..", "data", "audio")):
    os.mkdir(os.path.join("..", "..", "..", "data", "audio"))
    
if not os.path.isdir(os.path.join("..", "..", "..", "data", "audio_texts")):
    os.mkdir(os.path.join("..", "..", "..", "data", "audio_texts"))
    
os.environ["HF_HOME"] = os.path.join("..", "..", "..", "data", "hfcache")


transcriber = pipeline("automatic-speech-recognition", model='Simranjit/whisper-medical-french', device="cuda")

def transcribe(audio):

    sr, y = audio
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))


    text = transcriber({"sampling_rate": sr, "raw": y})["text"]
    text = text.replace("nouvelle ligne", "\n")
    text = text.replace("à la ligne", "\n")
    text = text.replace("  virgule", ",")
    text = text.replace(" virgule", ",")
    text = text.replace("  deux points", ":")
    text = text.replace(" deux points", ":")
    text = text.replace("  point", ".")
    text = text.replace(" point", ".")
    text = text.replace(" nouveau paragraphe ", "\n\n")
    text = text.replace(" paragraphe ", "\n\n")

    return text

def save_fn(audio, text):
    sr, y = audio
    y = y.astype(np.float32)
    y /= np.max(np.abs(y))

    uid = str(uuid.uuid4())

    with open(f"{uid}.txt", "w", encoding="utf-8") as f:
        f.write(text)
    s3.Bucket(BUCKET).upload_file(f"{uid}.txt", f"texts/{uid}.txt") #local path, bucket path

    write(f"{uid}.wav", sr, y)
    s3.Bucket(BUCKET).upload_file(f"{uid}.wav", f"audios/{uid}.wav") #local path, bucket path
    
    return [None, ""]

with gr.Blocks() as demo:
    audio = gr.Audio()
    text = gr.TextArea(show_copy_button=True)
    btn = gr.Button("run")
    btn.click(fn=transcribe, inputs=audio, outputs=text)
    save = gr.Button("save")
    save.click(fn=save_fn, inputs=[audio, text], outputs=[audio, text])

demo.launch(share=True)