from TTS.api import TTS
import numpy as np
import simpleaudio as sa
import torch
import threading
import queue

# Check if GPU is available
if torch.cuda.is_available():
    device = "cuda"
else:
    device = "cpu"

# Initialize the TTS object
tts = TTS(model_name="tts_models/multilingual/multi-dataset/xtts_v2", progress_bar=True)
tts.to(device)  # Use GPU if available

# Queue to hold audio chunks for playback
audio_queue = queue.Queue()

# Function to process text and generate audio chunks
def tts_worker(text, speaker="Ana Florence", language="en", chunk_size=10):
    words = text.split()
    chunks = [" ".join(words[i:i + chunk_size]) for i in range(0, len(words), chunk_size)]

    for chunk in chunks:
        print(f"Processing chunk: {chunk}")
        audio = tts.tts(
            text=chunk,
            speaker=speaker,
            language=language
        )
        audio_data = np.array(audio, dtype=np.float32)
        audio_data = (audio_data * 32767).astype(np.int16)  # Convert to 16-bit PCM
        audio_queue.put(audio_data)  # Add the audio chunk to the queue

    audio_queue.put(None)  # Signal the end of processing

# Function to play audio chunks
def audio_worker():
    while True:
        audio_data = audio_queue.get()  # Get the next audio chunk
        if audio_data is None:  # End of processing
            break
        play_obj = sa.play_buffer(audio_data, 1, 2, tts.synthesizer.output_sample_rate)
        play_obj.wait_done()  # Wait for the current chunk to finish playing

# Continuous loop for text input
print("Enter text to generate speech. Type 'exit' to quit.")
while True:
    # Get user input
    text = input("Enter text: ")

    # Exit the loop if the user types 'exit'
    if text.lower() == "exit":
        print("Exiting...")
        break

    # Start the TTS worker thread
    tts_thread = threading.Thread(target=tts_worker, args=(text, "Ana Florence", "en"))
    tts_thread.start()

    # Start the audio worker thread
    audio_thread = threading.Thread(target=audio_worker)
    audio_thread.start()

    # Wait for both threads to finish
    tts_thread.join()
    audio_thread.join()

    print("Streaming finished.")