#!/usr/local/lib/pytorch-venv/bin/python3
import os
from dotenv import load_dotenv
import gemini_srt_translator as gst


def translate(path, srt_file):
    load_dotenv()
    api_key = os.getenv('GEMINI_API_KEY')
    if not api_key:
        print("Error: Missing GEMINI_API_KEY environment variable.")
        return

    gst.gemini_api_key = api_key
    gst.target_language = "Persian"
    gst.input_file = os.path.join(path, srt_file)
    gst.output_file = os.path.join(path, 'translated', srt_file)
    gst.model_name = "gemini-2.0-flash"

    print(f"Translating {srt_file} ...")
    gst.translate()
    print("Translation completed successfully.")


if __name__ == '__main__':
    home_dir = os.path.expanduser('~')
    directory = os.path.join(home_dir, 'Documents', 'Subtitles')
    files = os.listdir(directory)
    srt_files = [f for f in files if f.endswith('.srt')]
    for file in srt_files:
        translate(directory, file)