import gradio as gr
import google.generativeai as genai
import re
import os
import time

# تنظیم API Key
API_KEY = os.getenv("GEMINI_API_KEY")  # کلیدت رو اینجا بذار یا از Secrets
genai.configure(api_key=API_KEY)
model = genai.GenerativeModel("gemini-1.5-flash")

def translate_with_gemini(text):
    if not text.strip():
        print("Text is empty, returning empty string")
        return ""
    print(f"Starting translation for: '{text}' at {time.ctime()}")
    start_time = time.time()
    try:
        prompt = f"Translate this to Persian: {text}"
        response = model.generate_content(prompt)
        translated = response.text
        elapsed_time = time.time() - start_time
        print(f"Translated to: '{translated}' (took {elapsed_time:.2f} seconds)")
        return translated
    except Exception as e:
        error_msg = f"خطا در ترجمه: {str(e)} (after {time.time() - start_time:.2f} seconds)"
        print(error_msg)
        return error_msg

def process_srt(file):
    encodings = ['utf-8', 'windows-1252', 'latin1']
    content = None
    for encoding in encodings:
        try:
            with open(file.name, 'r', encoding=encoding) as f:
                content = f.read()
            print(f"File read with encoding: {encoding}")
            break
        except UnicodeDecodeError:
            continue
    
    if content is None:
        raise ValueError("نمی‌تونم فایل رو با انکودینگ‌های رایج بخونم. لطفاً فایل رو با UTF-8 ذخیره کن.")

    blocks = re.split(r'\n\n', content.strip())
    print(f"Found {len(blocks)} blocks in SRT")
    subtitles = ["خروجی‌ها توسط هوش مصنوعی انجام شده\n"]
    for i, block in enumerate(blocks):
        print(f"Processing block {i + 1}/{len(blocks)}")
        lines = block.split('\n')
        if len(lines) >= 2:
            index = lines[0].strip()
            timing = lines[1].strip()
            text = '\n'.join(lines[2:]).strip()
            translated_text = translate_with_gemini(text)
            subtitles.append(f"{index}\n{timing}\n{translated_text}")
    output_content = '\n\n'.join(subtitles)
    output_path = "translated_srt.srt"
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(output_content)
    print("File written successfully")
    return output_path

interface = gr.Interface(
    fn=process_srt,
    inputs=gr.File(label="فایل SRT انگلیسی رو آپلود کن"),
    outputs=gr.File(label="فایل SRT ترجمه‌شده به فارسی"),
    title="ترجمه زیرنویس با Gemini API",
    description="فایل SRT انگلیسی رو آپلود کن تا به فارسی ترجمه بشه."
)

interface.launch()