import anyio
from argparse import ArgumentParser
from urllib.parse import urlparse, parse_qs
from youtube_transcript_api import YouTubeTranscriptApi
from multillm import Client

DEFAULT_MODEL = "claude/sonnet"
DEFAULT_SYSTEM_PROMPT = "Summarize the following video transcript:"


async def summarize(model: str, prompt: str):
    client = Client()
    result = await client.single(model, prompt)
    print(result)


def main():
    parser = ArgumentParser()
    parser.add_argument("url")
    parser.add_argument("-l", "--language", default="en", help="Language code for captions (e.g., en, es, fr)")
    parser.add_argument("-r", "--raw", action="store_true", help="Output raw transcript without summarizing")
    parser.add_argument("-m", "--model", default=DEFAULT_MODEL, help="Model (e.g., openai/gpt-4o, anthropic/claude-sonnet-4-20250514, claude/sonnet)")
    parser.add_argument("-s", "--system-prompt", default=DEFAULT_SYSTEM_PROMPT, help="Custom prompt for the LLM")
    args = parser.parse_args()
    api = YouTubeTranscriptApi()
    url = urlparse(args.url)
    qs = parse_qs(url.query)
    transcript = api.fetch(qs["v"][0], languages=[args.language])
    text = " ".join([s.text for s in transcript.snippets])
    if args.raw:
        print(text)
        return
    prompt = args.system_prompt + "\n" + text
    anyio.run(summarize, args.model, prompt)


if __name__ == "__main__":
    main()