import streamlit as st
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Load CodeT5 model from Hugging Face
MODEL_NAME = "Salesforce/codet5-large"
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_NAME)

def translate_code(code_snippet, source_lang, target_lang):
    """
    Translate code using CodeT5 model.
    """
    prompt = f"""Translate this {source_lang} code to {target_lang}:

{code_snippet}"""
    
    # Tokenize and generate translation
    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True, max_length=512)
    outputs = model.generate(**inputs, max_length=512)
    
    # Decode the output
    translated_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return translated_code

# Streamlit UI
st.title("🔄 AI Code Translator")
st.write("Convert code between Python, Java, C++, and C.")

languages = ["Python", "Java", "C++", "C"]
source_lang = st.selectbox("Select Source Language", languages)
target_lang = st.selectbox("Select Target Language", languages)
code_input = st.text_area("Enter your code:", height=200)

if st.button("Translate"):
    if source_lang == target_lang:
        st.warning("⚠️ Source and target languages must be different!")
    elif not code_input.strip():
        st.warning("⚠️ Please enter some code before translating.")
    else:
        with st.spinner("Translating... ⏳"):
            translated_code = translate_code(code_input, source_lang, target_lang)
            st.subheader(f"Translated {target_lang} Code:")
            st.code(translated_code, language=target_lang.lower())