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


def create_directory(directory_path):
    os.makedirs(directory_path, exist_ok=True)


def translate(input_dir, output_dir, srt_file, language):
    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 = language
    gst.input_file = os.path.join(input_dir, srt_file)
    gst.output_file = os.path.join(output_dir, srt_file)
    gst.model_name = "gemini-2.0-flash"

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


if __name__ == '__main__':
    home_dir = os.path.expanduser('~')
    default_input_dir = os.path.join(home_dir, 'Documents', 'Subtitles')
    default_output_dir = os.path.join(default_input_dir, 'translated')

    parser = argparse.ArgumentParser(
        description='Translate subtitles using GEMINI SRT translator'
    )
    parser.add_argument("-l",
                        "--language",
                        help="Target language for translation",
                        default="Persian",
                        choices=["English", "Persian", "German", "Italian", "Chinese", "Russian", "Turkish"])
    parser.add_argument("-i",
                        "--input",
                        help="Path to input SRT file",
                        default=default_input_dir)
    parser.add_argument("-o",
                        "--output",
                        help="Path to output translated SRT file",
                        default=default_output_dir)

    args = parser.parse_args()

    if args.input == default_input_dir and args.output == default_output_dir:
        create_directory(default_output_dir)
    files = os.listdir(args.input)
    srt_files = [f for f in files if f.endswith('.srt')]
    for file in srt_files:
        translate(args.input, args.output, file, args.language)

    print(f"Translated subtitles saved to {args.output}.")