import streamlit as st
from langchain.chains import LLMChain
from langchain_core.prompts import ChatPromptTemplate
from langchain_google_genai import ChatGoogleGenerativeAI
import os
import google.generativeai as genai

api_key = os.getenv("GEMINI_KEY2")

# Define the Chat Prompt Template
prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
            "You are a helpful assistant that translates {input_language} to {output_language}.",
        ),
        ("human", "{input}"),
    ]
)

# Initialize the Google Generative AI model
llm = ChatGoogleGenerativeAI(
    model="gemini-1.5-pro",
    temperature=0.5,
    max_tokens=None,
    timeout=None,
    max_retries=2,
    api_key=api_key 
)

# Define the chain for translation
translation_chain = LLMChain(prompt=prompt, llm=llm)

# Function to translate text
def translate_text(input_text, input_language, output_language):
    try:
        response = translation_chain({"input": input_text, "input_language": input_language, "output_language": output_language})
        translation = response["text"].strip()
    except Exception as e:
        st.error(f"Error translating text to {output_language}: {e}")
        translation = f"Sorry, we couldn't translate to {output_language} at the moment."
    
    return translation

# Streamlit UI
st.title("📝 Text Translation Bot with LangChain and Google Generative AI")

# User input
input_text = st.text_area("Enter text to translate:")

# Language choices
languages = ["Hindi", "Spanish", "German"]
output_language = st.selectbox("Select output language:", languages)

# Generate translation button
if st.button("Translate"):
    if input_text.strip():
        st.write("Translating text, please wait...")
        translated_text = translate_text(input_text, "English", output_language)
        st.write("### Translated Text")
        st.write(translated_text)
    else:
        st.warning("Please enter some text to translate.")