# ytsum, Youtube Summarizer, a tool to summarize the content of a Youtube video based on its audio, using OpenAI models.
import os

from openai import OpenAI 
from pydub import AudioSegment
from pytube import YouTube

def get_api_key():
    try:
        print("Retreiving OpenAI API Key...")

        OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')

        print(f"Retreived OpenAI API Key.")
        return OPENAI_API_KEY
    except Exception as e:
        print(f"\033[91mFailed to retrieve OpenAI API Key: {e}\033[0m")
        return None

# Downloads the audio from a YouTube video.
def download(video_url, output_filename='audio.mp4'):
    try:
        print("Downloading audio...")

        yt = YouTube(video_url)
        audio_stream = yt.streams.get_audio_only()
        audio_filename = audio_stream.download(filename=output_filename)

        print(f"Downloaded audio.")
        return audio_filename
    except Exception as e:
        print(f"\033[91mFailed to download audio: {e}\033[0m")
        return None

# Compresses an audio file to target size in MB.
def compress_audio(input_filename, target_size_mb=25):
    try:
        print("Compressing audio...")
        audio = AudioSegment.from_file(input_filename)

        # Split the input filename into directory and file components
        dir_name, file_name = os.path.split(input_filename)

        # Add 'compressed_' prefix to just the file name, not the entire path
        output_filename = os.path.join(dir_name, "compressed_" + file_name[:-4] + ".mp3")

        target_size_bits = target_size_mb * 8 * 1024 * 1024
        target_bitrate = int(target_size_bits / len(audio))

        compressed_audio = audio.set_frame_rate(24000).set_channels(1).set_sample_width(2)

        compressed_audio.export(output_filename, format="mp3", bitrate=str(target_bitrate) + "k")

        print(f"Compressed audio. New file: {output_filename}")
        return output_filename
    except Exception as e:
        print(f"Failed to compress audio: {e}")
        return None

# Transcribes the audio file using OpenAI's transcription service.
def transcription(OPENAI_API_KEY, audio_filename):
    try:
        print("Transcriptiting audio...")

        client = OpenAI(api_key=OPENAI_API_KEY)
        with open(audio_filename, "rb") as audio_file:
            transcript = client.audio.transcriptions.create(
            file=audio_file,
            model="whisper-1",
            )

        print("Transcripted audio.")
        return transcript
    except Exception as e:
        print(f"\033[91mFailed to transcript audio: {e}\033[0m")
        return None

# Generates a summary for a given transcript using GPT-4.
def summary(OPENAI_API_KEY, transcript):
    try:
        print("Summarizing transcript...")

        client = OpenAI(api_key=OPENAI_API_KEY)
        stream = client.chat.completions.create(
            model="gpt-4",
            messages=[
                {"role": "system", "content": "The prompt you will receive will be the transcript of a youtube video, your objective is to summarize the content of that transcript to the best of your ability, keep in mind there may be music or other parasite noises in the transcript."},
                {"role": "user", "content": str(transcript)}
                ],
            stream=True,
        )

        for chunk in stream:
            if chunk.choices[0].delta.content is not None:
                print(f"\033[92m{chunk.choices[0].delta.content}\033[0m", end="", flush=True)

        print("Summarized transcript.")
    except Exception as e:
        print(f"\033[91mFailed to summarize transcript: {e}\033[0m")
        return None

# Deletes the specified audio file.
def cleanup(audio_filename):
    try:
        print("Deleting audio...")

        os.remove(audio_filename)

        print(f"Deleted audio.")
    except Exception as e:
        print(f"\033[91mFailed to delete audio: {e}\033[0m")
        return None

# Main function to orchestrate the API Key retreival, download, transcription, summarization, and cleanup process.
def main():
    try:
        OPENAI_API_KEY = get_api_key()
        video_url = str(input("Video url: "))

        if OPENAI_API_KEY is not None:
            audio_filename = download(video_url)
            if audio_filename is not None:
                compressed_audio_filename = compress_audio(audio_filename)
                transcript = transcription(OPENAI_API_KEY, compressed_audio_filename)
                if transcript is not None:
                    summary(OPENAI_API_KEY, transcript)
                cleanup(compressed_audio_filename)
    except Exception as e:
        print(f"\033[91m{e}\033[0m")
    finally:
        input("Press Enter to exit...")

if __name__ == "__main__":
    main()
