import streamlit as st
import os
from transformers import pipeline
from langdetect import detect
from groq import Groq

# Load Hugging Face token from environment
HF_TOKEN = os.environ.get("homeo_doc")
if not HF_TOKEN:
    st.error("Missing Hugging Face API token. Set 'homeo_doc' in environment variables.")

# Initialize translation pipeline
translator = pipeline("translation", model="facebook/nllb-200-distilled-600M", token=HF_TOKEN)

# Language code mapping for NLLB-200
LANG_CODE_MAP = {
    'en': 'eng_Latn',    # English
    'ur': 'urd_Arab',    # Urdu
    'ar': 'arb_Arab',    # Arabic
    'es': 'spa_Latn',    # Spanish
    'hi': 'hin_Deva',    # Hindi
    'fr': 'fra_Latn'     # French
}

def translate_text(text, target_lang='eng_Latn'):
    """Translate text using NLLB-200"""
    try:
        source_lang = detect(text)
        source_code = LANG_CODE_MAP.get(source_lang, 'eng_Latn')  # Detect source language

        translation = translator(
            text, 
            src_lang=source_code,   # Pass source language
            tgt_lang=target_lang    # Pass target language
        )
        
        return translation[0]['translation_text']
    
    except Exception as e:
        st.error(f"Translation error: {str(e)}")
        return text

# Test function
if __name__ == "__main__":
    test_text = "یہ ایک آزمائشی جملہ ہے۔"  # Urdu sample text
    translated = translate_text(test_text, "eng_Latn")
    print("Translated:", translated)