import torch

import gradio as gr
from transformers import pipeline
from huggingface_hub import model_info
import Levenshtein


MODEL_NAME = "openai/whisper-small" #this always needs to stay in line 8 :D sorry for the hackiness
lang = "en"

device = 0 if torch.cuda.is_available() else "cpu"
pipe = pipeline(
    task="automatic-speech-recognition",
    model=MODEL_NAME,
    chunk_length_s=30,
    device=device,
)

pipe.model.config.forced_decoder_ids = pipe.tokenizer.get_decoder_prompt_ids(language=lang, task="transcribe")

def transcribe(microphone, file_upload):
    warn_output = ""
    if (microphone is not None) and (file_upload is not None):
        warn_output = (
            "WARNING: You've uploaded an audio file and used the microphone. "
            "The recorded file from the microphone will be used and the uploaded audio will be discarded.\n"
        )

    elif (microphone is None) and (file_upload is None):
        return "ERROR: You have to either use the microphone or upload an audio file"

    file = microphone if microphone is not None else file_upload

    text = pipe(file)["text"]
    def find_similar_words(target_word, word_list, threshold=3):
        similar_words = []
        for word in word_list:
            distance = Levenshtein.distance(target_word, word)
            if distance <= threshold:
                similar_words.append((word, distance))
        return similar_words
    word_list = ["OMRON digital thermomenter", "Warmax ginger drink with honey and lemon", "Panadol advance 500mg Tablet 96", "Panadol advance 500mg Table 24", "Panadol Actifast Tablet 20",
             "Adol 250mg suppository 10", "Adol 250mg Suspension 100ml", "Advil cold  & sinus Tablet 20", "Advil 200mg Tablet 24", "Advil 200mg Capsule 32"]
    
    target_word = text
    # Threshold for similarity (optional)
    threshold = 9

    similar_words = find_similar_words(target_word, word_list, threshold)

    if not similar_words:
        text = "No similar words found."
        print(text)
        return warn_output + text
    else:
        print("Similar words:")
        text = [f"{word}" for word, distance in similar_words][0]
        print(text[0])
        return warn_output + text
        # for word, distance in similar_words:
        #     print(f"{word} (Similarity score: {distance})")


    print('===========================================')

    print(text)


demo = gr.Blocks()

mf_transcribe = gr.Interface(
    fn=transcribe,
    inputs=[
        gr.inputs.Audio(source="microphone", type="filepath", optional=True),
        gr.inputs.Audio(source="upload", type="filepath", optional=True),
    ],
    outputs="text",
    layout="horizontal",
    theme="huggingface",
    title="Aumet: Naeem Project.",
    description=(
        "This Poject created for Naeem using Hugging face, transformers, torch and gradio"
    ),
    allow_flagging="never",
)


with demo:
    gr.TabbedInterface([mf_transcribe], ["Transcribe Audio"])

demo.launch(enable_queue=True)