from smolagents.tools import Tool
import pytubefix

class YouTubeTranscriptExtractor(Tool):
    description = "Extracts the transcript from a YouTube video."
    name = "youtube_transcript_extractor"
    inputs = {'video_url': {'type': 'string', 'description': 'The URL of the YouTube video.'}}
    output_type = "string"

    def forward(self, video_url: str) -> str:

        try:
          from pytubefix import YouTube
          # Create a YouTube object
          yt = YouTube(video_url)
          lang='en'

          # Get the video transcript
          if lang in yt.captions:
              transcript = yt.captions['en'].generate_srt_captions()

          else:
              transcript = yt.captions.all()[0].generate_srt_captions()
              lang=yt.captions.all()[0].code

          return lang + "transcript : " + transcript
        # return transcript


        except Exception as e:

          return f"An unexpected error occurred: {str(e)}"


    def __init__(self, *args, **kwargs):
        self.is_initialized = False