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

# Functie pentru extragerea si 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)
    except:
        print(f"Transcript not available for {video_id}")

# Functie pentru a împărți fișierele în bucăți de 3KB
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) > 3 * 1024:  # Dacă fișierul este mai mare de 3KB
                with open(full_path, 'r', encoding='utf-8') as f:
                    contents = f.read()

                part_num = 1
                for i in range(0, len(contents), 3 * 1024):  # Împărțim conținutul în bucăți de 3KB
                    part = contents[i:i+3*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 3KB, îl copiem pur și simplu în noul folder
                shutil.copy(full_path, split_folder)

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

    # Extragem ID-ul fiecarui 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')

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

        # După ce toate transcrierile au fost descărcate, împărțim fișierele
    split_files()

if __name__ == '__main__':
    main()