import gradio as gr
import numpy as np
from pydub import AudioSegment
import io
from transformers import pipeline, AutoTokenizer

# Load the pipeline for speech recognition and translation
pipe = pipeline(
    "automatic-speech-recognition",
    model="Akashpb13/Hausa_xlsr",
    tokenizer="Akashpb13/Hausa_xlsr"
)
translator = pipeline("text2text-generation", model="Baghdad99/saad-hausa-text-to-english-text")
tts = pipeline("text-to-speech", model="Baghdad99/english_voice_tts")

def translate_speech(audio_data_tuple):
    # Extract the audio data from the tuple
    sample_rate, audio_data = audio_data_tuple

    # Convert the audio data to int16 format
    audio_data_int16 = audio_data.astype(np.int16)

    # Create an AudioSegment from the audio data
    audio_segment = AudioSegment(
        audio_data_int16.tobytes(),  # Audio data as bytes
        frame_rate=sample_rate,
        sample_width=audio_data_int16.dtype.itemsize,  # Width in bytes
        channels=1
    )

    # Export the AudioSegment as MP3
    mp3_buffer = io.BytesIO()
    audio_segment.export(mp3_buffer, format="mp3")

    # Now you have an MP3 file in a BytesIO buffer. You can write it to a file,
    # send it over a network, etc. Here's how you can write it to a file:
    with open("audio.mp3", "wb") as f:
        f.write(mp3_buffer.getvalue())

    # Now you can feed the MP3 file to your model
    # Use the speech recognition pipeline to transcribe the audio
    output = pipe("audio.mp3")

    print(f"Output: {output}")  # Print the output to see what it contains

    # Check if the output contains 'text'
    if 'text' in output:
        transcription = output["text"]
    else:
        print("The output does not contain 'text'")
        return

    # Print the transcription
    print(f"Transcription: {transcription}")

    # Use the translation pipeline to translate the transcription
    translated_text = translator(transcription, return_tensors="pt")
    print(f"Translated text: {translated_text}")  # Print the translated text to see what it contains

    # Check if the translated text contains 'generated_token_ids'
    if 'generated_token_ids' in translated_text[0]:
        # Decode the tokens into text
        translated_text_str = translator.tokenizer.decode(translated_text[0]['generated_token_ids'])
    else:
        print("The translated text does not contain 'generated_token_ids'")
        return

    # Print the translated text string
    print(f"Translated text string: {translated_text_str}")

    # Use the text-to-speech pipeline to synthesize the translated text
    synthesised_speech = tts(translated_text_str)
    print(f"Synthesised speech: {synthesised_speech}")  # Print the synthesised speech to see what it contains

    # Check if the synthesised speech contains 'audio'
    if 'audio' in synthesised_speech:
        synthesised_speech_data = synthesised_speech['audio']
    else:
        print("The synthesised speech does not contain 'audio'")
        return

    # Flatten the audio data
    synthesised_speech_data = synthesised_speech_data.flatten()

    # Print the shape and type of the synthesised speech data
    print(f"Synthesised speech data type: {type(synthesised_speech_data)}, Synthesised speech data shape: {synthesised_speech_data.shape}")

    # Scale the audio data to the range of int16 format
    synthesised_speech = (synthesised_speech_data * 32767).astype(np.int16)

    return 16000, synthesised_speech


# Define the Gradio interface
iface = gr.Interface(
    fn=translate_speech, 
    inputs=gr.inputs.Audio(source="microphone"),  # Change this line
    outputs=gr.outputs.Audio(type="numpy"),
    title="Hausa to English Translation",
    description="Realtime demo for Hausa to English translation using speech recognition and text-to-speech synthesis."
)

iface.launch()