from huggingface_hub import InferenceClient
import gradio as gr
import re

client = InferenceClient("mistralai/Mixtral-8x7B-Instruct-v0.1")

def translate_subtitles(srt_content, target_language):
    """
    Translate the subtitles in the SRT file content to the target language.
    """
    # Split the SRT content into blocks
    blocks = srt_content.split('\n\n')
    translated_blocks = []
    for block in blocks:
        if block.strip() == "":
            continue
        lines = block.split('\n')
        if len(lines) >= 3:
            index = lines[0]
            time_range = lines[1]
            subtitle_text = '\n'.join(lines[2:])
            # Translate the subtitle text
            translation = client.translation(subtitle_text, target_language=target_language)
            translated_text = translation[0]['translation_text']
            translated_blocks.append(f"{index}\n{time_range}\n{translated_text}")
    return '\n\n'.join(translated_blocks)

def read_srt_file(file_path):
    """
    Read SRT file content.
    """
    with open(file_path, 'r', encoding='utf-8') as file:
        return file.read()

def save_translated_srt(content, output_path):
    """
    Save the translated subtitles to a new SRT file.
    """
    with open(output_path, 'w', encoding='utf-8') as file:
        file.write(content)

def translate_srt_interface(srt_file, target_language):
    """
    Gradio interface function to translate SRT file content.
    """
    srt_content = srt_file.read()
    translated_content = translate_subtitles(srt_content, target_language)
    return translated_content

iface = gr.Interface(
    fn=translate_srt_interface,
    inputs=[gr.File(file_type="srt"), gr.Textbox(label="Target Language Code")],
    outputs="text",
    title="SRT File Translator",
    description="Translate SRT subtitle files to your desired language."
)

iface.launch()