import streamlit as st
import pandas as pd
import pysrt
from transformers import MarianMTModel, MarianTokenizer

# Fetch and parse language options from the provided URL
url = "https://huggingface.co/Lenylvt/LanguageISO/resolve/main/iso.md"
df = pd.read_csv(url, delimiter="|", skiprows=2, header=None).dropna(axis=1, how='all')
df.columns = ['ISO 639-1', 'ISO 639-2', 'Language Name', 'Native Name']
df['ISO 639-1'] = df['ISO 639-1'].str.strip()

# Prepare language options for the dropdown
language_options = df.set_index('ISO 639-1')[['Language Name']].to_dict()['Language Name']

def translate_text(text, source_language_code, target_language_code):
    # Construct model name using ISO 639-1 codes
    model_name = f"Helsinki-NLP/opus-mt-{source_language_code}-{target_language_code}"

    # Check if source and target languages are the same, which is not supported for translation
    if source_language_code == target_language_code:
        return "Translation between the same languages is not supported."

    # Load tokenizer and model
    try:
        tokenizer = MarianTokenizer.from_pretrained(model_name)
        model = MarianMTModel.from_pretrained(model_name)
    except Exception as e:
        return f"Failed to load model for {source_language_code} to {target_language_code}: {str(e)}"

    # Translate text
    translated = model.generate(**tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512))
    translated_text = tokenizer.decode(translated[0], skip_special_tokens=True)
    
    return translated_text

def translate_srt(input_file, source_language_code, target_language_code):
    # Load SRT file
    subs = pysrt.open(input_file.name)

    # Initialize an empty list to store translated subtitles
    translated_subs = []

    # Translate each subtitle
    for idx, sub in enumerate(subs):
        translated_text = translate_text(sub.text, source_language_code, target_language_code)
        # Construct the translated subtitle with timestamp and line number
        translated_sub = pysrt.SubRipItem(index=idx+1, start=sub.start, end=sub.end, text=translated_text)
        translated_subs.append(translated_sub)

    # Save translated subtitles to a new SRT file
    translated_file = pysrt.SubRipFile(translated_subs)
    translated_srt_path = input_file.name.replace(".srt", f"_{target_language_code}.srt")
    translated_file.save(translated_srt_path)
    return translated_srt_path

st.title("SRT Translator")
st.write("Translate subtitles from one language to another.")

source_language_code = st.selectbox("Select Source Language", options=list(language_options.keys()))
target_language_code = st.selectbox("Select Target Language", options=list(language_options.keys()))

file_input = st.file_uploader("Upload SRT File", type=["srt"])

if file_input is not None:
    translated_srt_path = translate_srt(file_input, source_language_code, target_language_code)
    st.success(f"Translation complete! You can download the translated SRT file from [here]({translated_srt_path})")