import tempfile

from smolagents import Tool
import whisper
from pytubefix import YouTube
from pytubefix.cli import on_progress

class YoutubeTranscriberTool(Tool):
    name = "youtube_transcriber"
    description = "This is a tool that get the transcription of the YouTube video in the form of text."
    inputs = {
        "url": {
            "type": "string",
            "description": "The link of any youtube video to get the transcription",
        }
    }
    output_type = "string"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.model = whisper.load_model("base")


    def forward(self, url: str) -> str:
        yt = YouTube(url, on_progress_callback=on_progress)
        audio_stream = yt.streams.get_audio_only()
        temp_dir = tempfile.gettempdir()
        out_file = audio_stream.download(output_path=temp_dir)
        result = self.model.transcribe(out_file)
        return result['text']