import ffmpeg
import json
import os
import sys
import yt_dlp
from base64 import b64encode
from copy import deepcopy
from datetime import datetime
from decimal import *
from urllib.parse import urlparse, parse_qs


URL_CACHE_FILE = "stream_urls.json"


def get_stream_urls(video_url):
    ydl_opts = {
        "format": "best",
        "quiet": True,
    }
    with yt_dlp.YoutubeDL(ydl_opts) as ydl:
        info_dict = ydl.extract_info(video_url, download=False)
    for fmt in info_dict["formats"]:
        if fmt.get("protocol") == "https" and fmt.get("acodec") != "none":
            audio_url = fmt["url"]
        elif fmt.get("protocol") == "https" and fmt.get("vcodec") != "none":
            video_url = fmt["url"]
    return {
        "video": video_url,
        "audio": audio_url,
        "time": info_dict["release_timestamp"],
    }


def get_pyonmaku_tag(start_time, end_time, stream_timestamp):
    start_timestamp = int(stream_timestamp + start_time)
    end_timestamp = int(stream_timestamp + end_time)
    start_timestamp_base64 = (
        b64encode(start_timestamp.to_bytes(4)).decode("utf-8").rstrip("=")
    )
    end_timestamp_base64 = (
        b64encode(end_timestamp.to_bytes(4)).decode("utf-8").rstrip("=")
    )
    return f"!Pyonmaku({start_timestamp_base64}:{end_timestamp_base64})"


def check_if_expired(stream_url):
    url_expiry = float(parse_qs(urlparse(stream_url).query)["expire"][0])
    return datetime.now().timestamp() > url_expiry


def remove_old_urls(url_cache):
    pruned_cache = deepcopy(url_cache)
    for url, data in url_cache.items():
        if check_if_expired(data["video"]) or check_if_expired(data["audio"]):
            del pruned_cache[url]
    return pruned_cache


def main():
    url = sys.argv[1] if len(sys.argv) >= 2 else ""
    start_time = sys.argv[2] if len(sys.argv) >= 3 else ""
    end_time = sys.argv[3] if len(sys.argv) >= 4 else ""
    output_filename = sys.argv[4] if len(sys.argv) >= 5 else "clip.mp4"
    if not url or not start_time or not end_time:
        print(
            "Usage: clip_video.py <video url> <start time> <end time> [output filename]"
        )
        return
    start_time = sum(
        Decimal(x) * 60**i for i, x in enumerate(reversed(start_time.split(":")))
    )
    end_time = sum(
        Decimal(x) * 60**i for i, x in enumerate(reversed(end_time.split(":")))
    )
    url_cache = {}
    cache_file_path = os.path.join(
        os.path.dirname(os.path.abspath(sys.argv[0])), URL_CACHE_FILE
    )
    if os.path.exists(cache_file_path):
        with open(cache_file_path) as url_cache_file:
            url_cache = json.load(url_cache_file)
    if (
        url not in url_cache
        or check_if_expired(url_cache[url]["video"])
        or check_if_expired(url_cache[url]["audio"])
    ):
        url_cache = remove_old_urls(url_cache)
        url_cache[url] = get_stream_urls(url)
        with open(cache_file_path, "w") as url_cache_file:
            json.dump(url_cache, url_cache_file)
    video_input = ffmpeg.input(
        url_cache[url]["video"], ss=f"{start_time}", t=f"{end_time - start_time}"
    )
    audio_input = ffmpeg.input(
        url_cache[url]["audio"], ss=f"{start_time}", t=f"{end_time - start_time}"
    )
    ffmpeg.output(
        video_input,
        audio_input,
        output_filename,
        vcodec="libx264",
        crf="18",
        acodec="aac",
        audio_bitrate="192k",
    ).overwrite_output().run()
    pyonmaku_tag = get_pyonmaku_tag(start_time, end_time, url_cache[url]["time"])
    print(f"\nDone. Use {pyonmaku_tag} in description for Pyonmaku.")


if __name__ == "__main__":
    main()