import os
import streamlit as st
from typing import List

# Load Hugging Face Marian model (Helsinki-NLP) locally or via API

@st.cache_resource
def load_tokenizer_model(model_name: str):
    try:
        from transformers import MarianTokenizer, MarianMTModel
    except Exception:
        st.error("transformers library not installed. Run: pip install transformers sentencepiece")
        raise

    tokenizer = MarianTokenizer.from_pretrained(model_name)
    model = MarianMTModel.from_pretrained(model_name)
    return tokenizer, model


def translate_local(texts: List[str], tokenizer, model) -> List[str]:
    batch = tokenizer(texts, return_tensors="pt", padding=True, truncation=True)
    generated = model.generate(**batch, max_length=512)
    return tokenizer.batch_decode(generated, skip_special_tokens=True)


def translate_via_inference_api(text: str, model_name: str, hf_token: str) -> str:
    import requests
    API_URL = f"https://api-inference.huggingface.co/models/{model_name}"
    headers = {"Authorization": f"Bearer {hf_token}"}
    payload = {"inputs": text}
    res = requests.post(API_URL, headers=headers, json=payload, timeout=60)
    res.raise_for_status()
    data = res.json()
    if isinstance(data, dict) and data.get('error'):
        raise RuntimeError(data['error'])
    if isinstance(data, list) and len(data) > 0:
        return data[0].get('translation_text') or data[0].get('generated_text') or str(data[0])
    return str(data)

# ---------------- Streamlit UI ----------------
st.set_page_config(page_title="English → Urdu Translator", layout="centered")
st.title("🇬🇧 ➜ 🇵🇰 English → Urdu Translator")
st.write("Streamlit app that translates English text to Urdu using Hugging Face models.")

model_name = st.text_input("Model name", value="Helsinki-NLP/opus-mt-en-ur")
use_api = st.checkbox("Use Hugging Face Inference API instead of local model")

hf_token_env = os.environ.get("HF_API_TOKEN", "")
hf_token_input = st.text_input("Hugging Face API token (optional; required only for Inference API)", value=hf_token_env, type="password")

text_in = st.text_area("Input English text", height=200)

# Handle Streamlit version differences
try:
    col1, col2 = st.columns([1, 1])
except AttributeError:
    col1, col2 = st.beta_columns([1, 1])

with col1:
    translate_btn = st.button("Translate")
with col2:
    clear_btn = st.button("Clear")

if clear_btn:
    st.rerun()

if translate_btn:
    if not text_in.strip():
        st.warning("Please enter some English text to translate.")
    else:
        with st.spinner("Translating..."):
            try:
                if use_api:
                    hf_token = hf_token_input.strip()
                    if not hf_token:
                        st.error("Hugging Face API token required when using Inference API. Set as HF_API_TOKEN env var or paste it here.")
                    else:
                        result = translate_via_inference_api(text_in, model_name, hf_token)
                        st.subheader("Urdu translation")
                        st.write(result)
                else:
                    tokenizer, model = load_tokenizer_model(model_name)
                    outputs = translate_local([text_in], tokenizer, model)
                    st.subheader("Urdu translation")
                    st.write(outputs[0])
            except Exception as e:
                st.error(f"Translation failed: {e}")

st.markdown("---")
st.write("**Notes:**")
st.write("- Local model download may take some time on first run and requires internet.\n- To use the Hugging Face Inference API, enable the checkbox and provide your API token from https://huggingface.co/settings/tokens.\n- Install dependencies: `pip install streamlit transformers sentencepiece requests`.")

st.caption("Built with ❤️ — modify freely for your needs.")