import os

from smolagents import Tool
from supadata import Supadata


class SupadataTool(Tool):
    name = "supadata_tool"
    description = "Get the YouTube video URL and generate transcription describing the video content. If you need to know what the video is about, use this tool."
    inputs = {
        "url": {
            "type": "string",
            "description": "Full YouTube video URL",
        }
    }
    output_type = "string"

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.supadata = Supadata(api_key=os.getenv("SUPADATA_API_KEY"))

    def forward(self, url):
        video_id = self._extract_video_id(url)
        transcript = self.supadata.youtube.transcript(video_id=video_id, lang="en")
        return transcript.content

    @staticmethod
    def _extract_video_id(url):
        if not url.startswith('https://'):
            return url

        if 'youtube.com/watch?v=' in url:
            return url.split('watch?v=')[1].split('&')[0]
        elif 'youtu.be/' in url:
            return url.split('youtu.be/')[1].split('?')[0]
        return url