import gradio as gr
import numpy as np
import torch
import os
import tempfile
import logging
import asyncio
from typing import Optional
from transformers import AutoProcessor, AutoModelForSpeechSeq2Seq, AutoModelForTextToSpectrogram
from deep_translator import GoogleTranslator
import time
import threading
import torchaudio

# Set up logging
logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
logger = logging.getLogger(__name__)

# Device setup
device = "cuda:0" if torch.cuda.is_available() else "cpu"
logger.info(f"Device set to use {device}")

# Supported languages
SUPPORTED_LANGUAGES = {
    "English": "en",
    "Hindi": "hi",
}

# Load ASR (Speech Recognition) Model
asr_model_name = "openai/whisper-base"
asr_processor_name = "openai/whisper-base"

# Load TTS (Text-to-Speech) Model
tts_model_name = "microsoft/speecht5_tts"
tts_processor_name = "microsoft/speecht5_tts"

if 'SPACE_ID' in os.environ:  # Check if running on Hugging Face Spaces
    cache_dir = "./models"
    asr_processor = AutoProcessor.from_pretrained(asr_processor_name, cache_dir=cache_dir)
    asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(asr_model_name, cache_dir=cache_dir).to(device)
    tts_processor = AutoProcessor.from_pretrained(tts_processor_name, cache_dir=cache_dir)
    tts_model = AutoModelForTextToSpectrogram.from_pretrained(tts_model_name, cache_dir=cache_dir).to(device)
else:
    asr_processor = AutoProcessor.from_pretrained(asr_processor_name)
    asr_model = AutoModelForSpeechSeq2Seq.from_pretrained(asr_model_name).to(device)
    tts_processor = AutoProcessor.from_pretrained(tts_processor_name)
    tts_model = AutoModelForTextToSpectrogram.from_pretrained(tts_model_name).to(device)

def speech_to_text(audio_data: np.ndarray, sample_rate: int) -> str:
    try:
        if len(audio_data.shape) > 1:
            audio_data = np.mean(audio_data, axis=1)
        audio_data = torch.tensor(audio_data, dtype=torch.float32).to(device)
        
        inputs = asr_processor(audio_data, sampling_rate=sample_rate, return_tensors="pt").to(device)
        outputs = asr_model.generate(**inputs)
        text = asr_processor.batch_decode(outputs, skip_special_tokens=True)[0]
        logger.info(f"Transcribed text: {text}")
        return text if text else ""
    except Exception as e:
        logger.error(f"STT failed: {str(e)}")
        return ""

def text_to_speech(text: str, lang: str) -> Optional[str]:
    try:
        if not text or not text.strip():
            logger.warning("Empty text received for TTS.")
            return None
        target_lang_code = SUPPORTED_LANGUAGES.get(lang, "en")
        translated_text = GoogleTranslator(source="auto", target=target_lang_code).translate(text)
        
        inputs = tts_processor(translated_text, return_tensors="pt").to(device)
        spectrogram = tts_model.generate(**inputs)
        
        with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as temp_file:
            torchaudio.save(temp_file.name, spectrogram.cpu(), 16000)
            return temp_file.name
    except Exception as e:
        logger.error(f"TTS failed: {str(e)}")
        return None

async def process_audio_stream(audio_stream, target_lang: str):
    async for sample_rate, audio_chunk in audio_stream:
        try:
            if audio_chunk is None or not isinstance(audio_chunk, np.ndarray):
                yield "", None
                continue
            text = speech_to_text(audio_chunk, sample_rate)
            dubbed_audio = text_to_speech(text, target_lang) if text and text.strip() else None
            yield text, dubbed_audio
        except Exception as e:
            logger.error(f"Stream processing failed: {str(e)}")
            yield "", None

def run_cleanup():
    while True:
        for root, dirs, files in os.walk(tempfile.gettempdir()):
            for file in files:
                if file.endswith(".wav"):
                    try:
                        os.remove(os.path.join(root, file))
                    except:
                        pass
        time.sleep(300)

with gr.Blocks(title="Real-Time Multilingual Dubbing") as demo:
    gr.Markdown("<h1>Real-Time Multilingual Dubbing</h1>")
    with gr.Row():
        audio_input = gr.Audio(sources=["microphone"], type="numpy", streaming=True, label="Speak")
        lang_dropdown = gr.Dropdown(choices=list(SUPPORTED_LANGUAGES.keys()), label="Target Language", value="English")
    with gr.Row():
        stt_output = gr.Textbox(label="Transcription")
        dub_output = gr.Audio(label="Dubbed Audio", autoplay=True)
    
    audio_input.stream(fn=process_audio_stream, inputs=[audio_input, lang_dropdown], outputs=[stt_output, dub_output])

if __name__ == "__main__":
    cleanup_thread = threading.Thread(target=run_cleanup, daemon=True)
    cleanup_thread.start()
    demo.launch()