import streamlit as st
import pandas as pd
from transformers import pipeline
from datetime import datetime

# ================================
# Cache the translation pipelines
# ================================
@st.cache_resource
def load_translation_pipelines():
    """
    Load and cache translation pipelines to avoid reloading on every interaction.
    """
    try:
        enja = pipeline("translation", model="staka/fugumt-en-ja")
        jaen = pipeline("translation", model="staka/fugumt-ja-en")
        zhja = pipeline("translation", model="Helsinki-NLP/opus-mt-tc-big-zh-ja")
        return {'enja': enja, 'jaen': jaen, 'zhja': zhja}
    except Exception as e:
        st.error(f"Error loading translation models: {e}")
        return {}

# Load the translation models
session_models = load_translation_pipelines()

# ================================
# Streamlit Application Layout
# ================================
st.set_page_config(
    page_title="Multi-Language Translator",
    layout="centered",
    initial_sidebar_state="auto",
)

st.title("🌐 Multi-Language Translator")

# Initialize session state for CSV creation flag
if 'csv_created' not in st.session_state:
    st.session_state.csv_created = False

# ================================
# User Input Section
# ================================
st.header("🔤 Enter Text to Translate")

# Model selection
model_options = {
    'English to Japanese': 'enja',
    'Japanese to English': 'jaen',
    'Chinese to Japanese': 'zhja'
}
model_display = list(model_options.keys())
model_keys = list(model_options.values())

selected_model_display = st.selectbox("Select Translation Model", model_display, index=0)
selected_model = model_options[selected_model_display]

# Text input
text = st.text_area("Input Text", height=150)

# ================================
# Translation and Output
# ================================
if st.button("🚀 Translate"):
    if not text.strip():
        st.warning("Please enter text to translate.")
    elif selected_model not in session_models:
        st.error("Selected translation model is not available.")
    else:
        with st.spinner("Translating..."):
            try:
                translator = session_models[selected_model]
                translation = translator(text)[0]['translation_text']
                st.success("Translation Successful!")
                st.subheader("📝 Translation Result")
                st.write(translation)

                # Prepare data for CSV
                timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                data = {
                    'Timestamp': [timestamp],
                    'Model': [selected_model_display],
                    'Original Text': [text],
                    'Translated Text': [translation]
                }
                df = pd.DataFrame(data)

                # Save to CSV
                csv_file = 'translation_data.csv'
                if not st.session_state.csv_created:
                    df.to_csv(csv_file, mode='w', header=True, index=False)
                    st.session_state.csv_created = True
                else:
                    df.to_csv(csv_file, mode='a', header=False, index=False)

                st.info(f"Translation saved to `{csv_file}`.")
            except Exception as e:
                st.error(f"An error occurred during translation: {e}")

# ================================
# Optional: Download Translation Data
# ================================
if st.button("📥 Download Translation Data"):
    try:
        df = pd.read_csv('translation_data.csv')
        csv = df.to_csv(index=False).encode('utf-8')
        st.download_button(
            label="Download CSV",
            data=csv,
            file_name='translation_data.csv',
            mime='text/csv',
        )
    except FileNotFoundError:
        st.warning("No translation data available to download.")
    except Exception as e:
        st.error(f"An error occurred while preparing the download: {e}")