import os
import shutil
from youtube_transcript_api import YouTubeTranscriptApi
from pytube import YouTube
from concurrent.futures import ProcessPoolExecutor
from tqdm import tqdm

# Funcție pentru extragerea și salvarea transcrierilor
def extract_transcript(video_id):
    try:
        transcript = YouTubeTranscriptApi.get_transcript(video_id)
        transcript_text = ' '.join([t['text'] for t in transcript])

        # Obținem titlul videoclipului
        video = YouTube(f"https://www.youtube.com/watch?v={video_id}")
        title = video.title

        # Curățăm titlul pentru a-l face compatibil cu numele fișierelor
        clean_title = "".join(c for c in title if c.isalnum() or c.isspace())

        with open(f'Transcripts/{clean_title}.txt', 'w', encoding='utf-8') as out_file:
            out_file.write(transcript_text)

        # Traducerea transcrierii în limba română
        translated_transcript = transcript.translate('ro')
        translated_text = translated_transcript.fetch()

        # Salvarea transcrierii traduse într-un folder separat
        with open(f'Transcripts_Romanian/{clean_title}_ro.txt', 'w', encoding='utf-8') as out_file:
            out_file.write(translated_text)
    except:
        print(f"Transcript not available for {video_id}")

# Funcție pentru împărțirea fișierelor în bucăți de 300KB
def split_files():
    split_folder = 'Transcripts/data_split'

    if not os.path.exists(split_folder):
        os.makedirs(split_folder)

    for filename in os.listdir('Transcripts'):
        if filename.endswith('.txt'):
            full_path = f'Transcripts/{filename}'
            if os.path.getsize(full_path) > 300 * 1024:  # Dacă fișierul este mai mare de 300KB
                with open(full_path, 'r', encoding='utf-8') as f:
                    contents = f.read()

                part_num = 1
                for i in range(0, len(contents), 300 * 1024):  # Împărțim conținutul în bucăți de 300KB
                    part = contents[i:i+300*1024]
                    with open(f'{split_folder}/{filename[:-4]}_part_{part_num}.txt', 'w', encoding='utf-8') as out_file:
                        out_file.write(part)

                    part_num += 1
            else:
                # Dacă fișierul este mai mic de 300KB, îl copiem pur și simplu în noul folder
                shutil.copy(full_path, split_folder)

def main():
    # Citirea linkurilor din fișierul logs.txt
    with open('logs.txt', 'r') as f:
        video_links = f.read().splitlines()

    # Extragem ID-ul fiecărui videoclip din link
    video_ids = [link.split('v=')[-1] for link in video_links]

    # Crearea unui folder pentru a stoca transcrierile
    if not os.path.exists('Transcripts'):
        os.makedirs('Transcripts')

    # Crearea unui folder pentru a stoca transcrierile traduse

        if not os.path.exists('Transcripts_Romanian'):
        os.makedirs('Transcripts_Romanian')

    # Utilizarea unui pool de procese pentru a executa extragerile în paralel
    with ProcessPoolExecutor() as executor:
        list(tqdm(executor.map(extract_transcript, video_ids), total=len(video_ids)))

    # După ce toate transcrierile sunt extrase și salvate, împărțim fișierele în bucăți
    split_files()

if __name__ == '__main__':
    main()