from transformers import MarianMTModel, MarianTokenizer, pipeline
import gradio as gr
import tempfile
from gtts import gTTS

MODEL_NAME = "victorachede/tiv-translator"

print("Loading tokenizer...")
tokenizer = MarianTokenizer.from_pretrained(MODEL_NAME)
print("Loading model...")
model = MarianMTModel.from_pretrained(MODEL_NAME)
print("Loading Whisper...")
asr = pipeline("automatic-speech-recognition", model="openai/whisper-base")
print("All models ready.")

def translate(text: str) -> str:
    if not text or not text.strip():
        return ""
    inputs = tokenizer(
        text.strip(),
        return_tensors="pt",
        padding=True,
        truncation=True,
        max_length=512
    )
    outputs = model.generate(
        **inputs,
        max_length=128,
        num_beams=5,
        repetition_penalty=1.3,
        no_repeat_ngram_size=3,
        early_stopping=True
    )
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

def text_to_speech(tiv_text: str) -> str:
    tts = gTTS(text=tiv_text, lang='en')  # placeholder — swap for ElevenLabs later
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
    tts.save(tmp.name)
    return tmp.name

def translate_text(english_text: str):
    tiv = translate(english_text)
    audio = text_to_speech(tiv)
    return tiv, audio

def speech_to_speech(audio_path: str):
    result = asr(audio_path)
    english_text = result["text"]
    tiv = translate(english_text)
    audio = text_to_speech(tiv)
    return english_text, tiv, audio

with gr.Blocks(title="TRANSLTR — by Black Sheep Co.") as demo:
    gr.Markdown("# TRANSLTR\n### English → Tiv | by Black Sheep Co.")

    with gr.Tab("Text"):
        text_in = gr.Textbox(label="English", placeholder="Enter English text...", lines=3)
        text_out = gr.Textbox(label="Tiv", lines=3)
        audio_out_text = gr.Audio(label="Tiv (spoken)")
        text_in.submit(translate_text, inputs=text_in, outputs=[text_out, audio_out_text])
        gr.Button("Translate").click(translate_text, inputs=text_in, outputs=[text_out, audio_out_text])

    with gr.Tab("Speech"):
        audio_in = gr.Audio(sources=["microphone"], type="filepath", label="Speak English")
        english_heard = gr.Textbox(label="English (heard)")
        tiv_out = gr.Textbox(label="Tiv (translated)")
        audio_out_speech = gr.Audio(label="Tiv (spoken)")
        gr.Button("Translate").click(speech_to_speech, inputs=audio_in, outputs=[english_heard, tiv_out, audio_out_speech])

demo.launch()