import time
import requests
import re
import gradio as gr
import os
import shutil
api_key = "268976:66f4f58a2a905"


def read_srt_file(file_path):
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            srt_content = file.read()
            return srt_content
    except FileNotFoundError:
        print(f"The file {file_path} was not found.")
    except Exception as e:
        print(f"An error occurred: {e}")

def clean_text(text):
    # Remove 'srt ' from the start of each line
    # Remove ''' from the start and end
    text = re.sub(r"^```|```$", '', text)
    text = re.sub(r'^srt', '', text, flags=re.MULTILINE)
    return text

def translate_text(api_key, text, source_language = "en", target_language = "fa"):
    url = "https://api.one-api.ir/translate/v1/google/"
    request_body = {"source": source_language, "target": target_language, "text": text}
    headers = {"one-api-token": api_key, "Content-Type": "application/json"}
    response = requests.post(url, headers=headers, json=request_body)
    if response.status_code == 200:
        result = response.json()
        return result['result']
    else:
        print(f"Error: {response.status_code}, {response.text}")
        return None

def enhance_text(api_key, text):
    url = "https://api.one-api.ir/chatbot/v1/gpt4o/"

    # Prepare the request body
    request_body = [{
        "role": "user",
        "content":  f"Please take the following SRT subtitle text in English and translate only the subtitle text into Persian. Ensure that all numbering and time codes remain unchanged. convert English terms in to common persian terms. The output should be a new SRT file with the subtitles in Persian, preserving the original formatting and timings and exept for the subtitle dont return anything in response. the subtitle will be provided in the following message"
    },
    {
    "role": "assistant",
    "content": "okay"
    },
    {
    "role": "user",
    "content": text
    }
    ]

    # Add the API key to the request
    headers = {
        "one-api-token": api_key,
        "Content-Type": "application/json"
    }

    # Make the POST request
    attempts = 0
    max_attempts = 3

    while attempts < max_attempts:
        response = requests.post(url, headers=headers, json=request_body)
        if response.status_code == 200:
            result = response.json()
            if result["status"] == 200:
                print("status: ", result["status"])
                te = clean_text(result["result"][0])
                print("result: ", te)
                return te
            else:
                print(f"Error: status {result['status']}, retrying in 30 seconds...")
        else:
            print(f"Error: {response.status_code}, {response.text}, retrying in 30 seconds...")
        attempts += 1
        time.sleep(10)
    print("Error Max attempts reached. Could not retrieve a successful response.")
    te = translate_text(api_key, text)
    return te


def generate_translated_subtitle(language, segments, input_video_name):
    input_video_name=input_video_name.split('/')[-1]
    subtitle_file = f"{input_video_name}.srt"
    text = ""
    lines = segments.split('\n')
    new_list = [item for item in lines if item != '']
    segment_number = 1

    for index, segment in enumerate(new_list):
        if (index+1) % 3 == 1 or (index+1)==1:
            text += f"{segment}\n"
            segment_number += 1
        if (index+1) % 3 == 2 or (index+1)==2:
            text += segment + "\n"
        if (index+1) % 3 == 0:
            text += f"\u200F{segment}\n\n"

    with open(subtitle_file, "a", encoding='utf8') as f:
        f.write(text)
    return subtitle_file

def split_srt_file(input_file, max_chars=3000):
    # Read the contents of the SRT file
    with open(input_file, 'r', encoding='utf-8') as file:
        content = file.read()
    file.close()

    # Split the content into individual subtitles
    subtitles = content.strip().split('\n\n')

    # Prepare to write the split files
    output_files = []
    current_file_content = ''
    current_file_index = 1

    for subtitle in subtitles:
        # Check if adding this subtitle would exceed the character limit
        if len(current_file_content) + len(subtitle) + 2 > max_chars:  # +2 for \n\n
            # Write the current file
            output_file_name = f'split_{current_file_index}.srt'
            with open(output_file_name, 'w', encoding='utf-8') as output_file:
                output_file.write(current_file_content.strip())
            output_files.append(output_file_name)

            # Prepare for the next file
            current_file_index += 1
            current_file_content = subtitle + '\n\n'
        else:
            # If it fits, add the subtitle
            current_file_content += subtitle + '\n\n'

    # Write any remaining content to a new SRT file
    if current_file_content:
        output_file_name = f'split_{current_file_index}.srt'
        with open(output_file_name, 'w', encoding='utf-8') as output_file:
            output_file.write(current_file_content.strip())
        output_files.append(output_file_name)

    return output_files

def translate(file, max_chars):
    print("translate")
    srt_files = split_srt_file(file, max_chars=3000)
    for srt_file in srt_files:
        srt = read_srt_file(srt_file)
        srt_string = enhance_text(api_key, srt)
        print(srt_string)
        subtitle_file = generate_translated_subtitle('fa', srt_string, 'video_subtitled')
        time.sleep(10)

    if os.path.exists(subtitle_file):
        copied_subtitle = os.path.join(os.path.dirname(subtitle_file), 
                                       f"copy_{os.path.basename(subtitle_file)}")
        shutil.copy2(subtitle_file, copied_subtitle)
        
        # Delete the original subtitle file
        os.remove(subtitle_file)
        
        return copied_subtitle
    else:
        print("Error: subtitle file not found.")
        return None

with gr.Blocks() as demo:
    gr.Markdown("Start typing below and then click **Run** to see the progress and final output.")
    with gr.Column():
        srt_file = gr.File()
        max_chars = gr.Number()
        subtitle_file = gr.File()
        btn = gr.Button("Create")
        btn.click(
            fn=translate,
            inputs=[srt_file, max_chars],
            outputs=subtitle_file,
        )

demo.launch(debug=True)