import gradio as gr
import wave
import numpy as np
from io import BytesIO
from huggingface_hub import hf_hub_download
from piper import PiperVoice 
from transformers import pipeline

# Load the NSFW classifier model
nsfw_detector = pipeline("text-classification", model="michellejieli/NSFW_text_classifier")

def synthesize_speech(text):
    # Check for NSFW content
    nsfw_result = nsfw_detector(text)
    if nsfw_result[0]['label'] == 'NSFW':
        return "NSFW content detected. Cannot process.", None

    model_path = hf_hub_download(repo_id="DLI-SLQ/speaker_01234", filename="speaker__01234_model.onnx")
    config_path = hf_hub_download(repo_id="DLI-SLQ/speaker_01234", filename="speaker__01234_model.onnx.json")
    voice = PiperVoice.load(model_path, config_path)

    # Create an in-memory buffer for the WAV file
    buffer = BytesIO()
    with wave.open(buffer, 'wb') as wav_file:
        wav_file.setframerate(voice.config.sample_rate)
        wav_file.setsampwidth(2)  # 16-bit
        wav_file.setnchannels(1)  # mono

        # Synthesize speech
        voice.synthesize(text, wav_file)

    # Convert buffer to NumPy array for Gradio output
    buffer.seek(0)
    audio_data = np.frombuffer(buffer.read(), dtype=np.int16)
    return audio_data.tobytes(), None
    

# Gradio Interface
with gr.Blocks(theme=gr.themes.Base()) as blocks:
    gr.Markdown("# Text to Speech Synthesizer")
    gr.Markdown("Enter text to synthesize it into speech using models from the State Library of Queensland's collection using Piper.")
    input_text = gr.Textbox(label="Input Text")
    output_audio = gr.Audio(label="Synthesized Speech", type="numpy")
    output_text = gr.Textbox(label="Output Text", visible=True)  # Make this visible for error messages
    submit_button = gr.Button("Synthesize")

    def process_and_output(text):
        audio, message = synthesize_speech(text)
        if message:
            return None, message  # Return None for audio and the error message
        else:
            return audio, None  # Return the audio data and None for the message

    submit_button.click(process_and_output, inputs=input_text, outputs=[output_audio, output_text])

# Run the app
blocks.launch()