import os
import asyncio
import torch
import streamlit as st
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
from googletrans import Translator
import langdetect

# ✅ Disable file watching for PyTorch compatibility with Streamlit
os.environ["STREAMLIT_WATCH_FILE_SYSTEM"] = "false"

# ✅ Fix for asyncio event loop crash on Windows
try:
    asyncio.get_running_loop()
except RuntimeError:
    asyncio.set_event_loop(asyncio.new_event_loop())

# ✅ Load Fine-Tuned Chhattisgarhi Translation Model
model_path = "chhattisgarhi_translator"
model = MBartForConditionalGeneration.from_pretrained(model_path)
tokenizer = MBart50TokenizerFast.from_pretrained(model_path, src_lang="hi_IN", tgt_lang="hne_IN")
translator = Translator()

# ✅ Detect Language
def detect_language(text):
    try:
        return langdetect.detect(text)
    except:
        return "unknown"

# ✅ Translate English → Hindi
def translate_english_to_hindi(text):
    translated = translator.translate(text, src="en", dest="hi")
    return translated.text

# ✅ Translate Hindi → Chhattisgarhi
def translate_hindi_to_chhattisgarhi(text):
    sentences = text.split("।")  # Sentence splitting
    translated_sentences = []

    for sentence in sentences:
        sentence = sentence.strip()
        if sentence:
            inputs = tokenizer(sentence, return_tensors="pt", truncation=True, padding="longest", max_length=256)
            with torch.no_grad():
                translated_ids = model.generate(**inputs, max_length=256, num_beams=5, early_stopping=True)
            translated_text = tokenizer.decode(translated_ids[0], skip_special_tokens=True)
            translated_sentences.append(translated_text)

    return " । ".join(translated_sentences)

# ✅ Streamlit UI
st.title("English/Hindi to Chhattisgarhi Translator 🗣️")
st.write("Enter an English or Hindi sentence and get its translation in Chhattisgarhi.")

user_input = st.text_area("Enter text:")

if st.button("Translate"):
    if user_input.strip():
        lang = detect_language(user_input)

        if lang == "en":
            hindi_text = translate_english_to_hindi(user_input)
            chhattisgarhi_text = translate_hindi_to_chhattisgarhi(hindi_text)
            st.success(f"**Hindi Translation**:\n{hindi_text}")
            st.success(f"**Chhattisgarhi Translation**:\n{chhattisgarhi_text}")

        elif lang == "hi":
            chhattisgarhi_text = translate_hindi_to_chhattisgarhi(user_input)
            st.success(f"**Chhattisgarhi Translation**:\n{chhattisgarhi_text}")

        else:
            st.error("❌ Unable to detect language. Please enter text in English or Hindi.")
    else:
        st.warning("⚠ Please enter some text before translating.")