import os
import re
from pytube import YouTube
from youtube_transcript_api import YouTubeTranscriptApi
from transformers import MarianMTModel, MarianTokenizer

def download_video_transcript(video_url, source_lang='en', target_lang='fr'):
    try:
        video_id = re.search(r"(?<=v=)[\w-]+", video_url)
        if video_id:
            video_id = video_id.group()
            yt = YouTube(video_url)
            stream = yt.streams.get_highest_resolution()

            print(f'Downloading {yt.title}...')
            modified_title = yt.title.replace(" ", "_")
            download_path = 'CS370_Milestone5/videos'  # Updated path to CS370_Milestone5 repository
            captions_path = 'CS370_Milestone5/captions'  # Updated path to CS370_Milestone5 repository

            # Create directories if they don't exist
            os.makedirs(download_path, exist_ok=True)
            os.makedirs(captions_path, exist_ok=True)

            video_file = f'{download_path}/{modified_title}.mp4'
            stream.download(output_path=download_path, filename=modified_title + '.mp4')
            print('Download completed!')

            transcript_list = YouTubeTranscriptApi.list_transcripts(video_id)
            transcript = transcript_list.find_generated_transcript([source_lang]).fetch()

            original_captions = ""
            for i, line in enumerate(transcript):
                start_time = line['start']
                formatted_time = f"{int(start_time // 60):02d}:{int(start_time % 60):02d}"
                original_captions += f"{formatted_time} {line['text']}\n"

            original_filename = f'{captions_path}/{modified_title}_original.txt'
            with open(original_filename, 'w', encoding='utf-8') as file:
                file.write(original_captions)
            
            # Translation part
            translated_captions = translate_text_file(original_filename,
                                                      f'{captions_path}/{modified_title}_translated.txt',
                                                      source_lang=source_lang,
                                                      target_lang=target_lang)

            # Delete video file after captions are saved
            os.remove(video_file)
            
            return original_captions, translated_captions, original_filename, f'{captions_path}/{modified_title}_translated.txt'
        else:
            print("Video ID not found in URL.")
            return None, None, None, None
    except Exception as e:
        print("Error:", e)
        return None, None, None, None