#!/usr/bin/env python3

##########################################################################################
# This script is used to import subscriptions and playlists from Piped to Invidious.
#
# It requires the following:
#
#   - python3 + requests
#   - The Piped API URL (e.g. https://pipedapi.example.com).
#     Note: This is NOT the Piped frontend URL.
#   - The Invidious URL (e.g. https://invidious.example.com).
#   - A Piped API token (it can be obtained from an authenticated browser session),
#     passed on the `PIPED_TOKEN` environment variable
#   - The Invidious cookie header passed on the `INVIDIOUS_TOKEN` environment variable.
#     You can retrieve it in your browser via https://<invidious-url>/authorize_token?scopes=:*.
#
# Author: Fabio Manganiello <fabio@manganiello.tech>
# License: MIT
##########################################################################################

import argparse
import os
import re
import sys

from abc import ABC, abstractmethod
from typing import Sequence
from urllib.parse import urljoin

import requests


class Session:
    """
    A session contains the data required to interact with the Piped API and the Invidious instance.
    """

    def __init__(
        self,
        piped_api_url: str,
        invidious_url: str,
        piped_token: str,
        invidious_token: str,
        skip_playlists: bool = False,
        skip_subscriptions: bool = False,
    ):
        self.piped_api_url = piped_api_url
        self.invidious_url = invidious_url
        self.piped_token = piped_token
        self.invidious_token = invidious_token
        self.skip_playlists = skip_playlists
        self.skip_subscriptions = skip_subscriptions
        self.piped = PipedAPI(self)
        self.invidious = InvidiousAPI(self)

    @classmethod
    def from_cmdline(cls, args: Sequence[str]) -> "Session":
        argparser = argparse.ArgumentParser(
            description="Import subscriptions and playlists from Piped to Invidious."
        )

        argparser.add_argument(
            "--piped-api-url",
            required=True,
            help="The base URL of the Piped API of the instance you want to import from, e.g. https://pipedapi.example.com",
        )

        argparser.add_argument(
            "--invidious-url",
            required=True,
            help="The base URL of the Invidious instance you want to import to, e.g. https://invidious.example.com",
        )

        argparser.add_argument(
            "--skip-playlists", action="store_true", help="Skip importing playlists"
        )

        argparser.add_argument(
            "--skip-subscriptions",
            action="store_true",
            help="Skip importing subscriptions",
        )

        piped_token = os.getenv("PIPED_TOKEN")
        invidious_token = os.getenv("INVIDIOUS_TOKEN")

        if not piped_token:
            argparser.error("The PIPED_TOKEN environment variable is required")

        if not invidious_token:
            argparser.error(
                "The INVIDIOUS_TOKEN environment variable is required.\n"
                "You can retrieve it in your browser via https://<invidious-url>/authorize_token?scopes=:*"
            )

        opts, _ = argparser.parse_known_args(args)
        return cls(
            piped_api_url=opts.piped_api_url,
            invidious_url=opts.invidious_url,
            piped_token=piped_token,
            invidious_token=invidious_token,
            skip_playlists=opts.skip_playlists,
            skip_subscriptions=opts.skip_subscriptions,
        )

    def migrate_subscriptions(self):
        if self.skip_subscriptions:
            print("=== Skipping subscriptions import ===")
            return

        print("=== Importing subscriptions ===")

        try:
            rs = self.piped.request("/subscriptions")
        except requests.exceptions.HTTPError as e:
            sys.stderr.write(f"===> [-] Could not fetch Piped subscriptions: {e}\n")
            sys.exit(1)

        piped_subs = rs.json()
        print(f"=== Found {len(piped_subs)} Piped subscriptions ===")

        invidious_subs = {
            sub["authorId"]
            for sub in self.invidious.request("/auth/subscriptions").json()
        }

        for subs in piped_subs:
            channel = re.sub(r"^/channel/", "", subs["url"])
            if channel in invidious_subs:
                print(
                    f'[=] Channel "{subs["name"]}" already subscribed on Invidious, skipping'
                )
                continue

            print(f'[+] Importing channel: {subs["name"]}')

            try:
                rs = self.invidious.request(
                    f"/auth/subscriptions/{channel}", method="POST"
                )
            except requests.exceptions.HTTPError as e:
                sys.stderr.write(
                    f'===> [-] Could not subscribe to channel: {subs["name"]}: {e}\n'
                )

    def migrate_playlists(self):
        if self.skip_playlists:
            print("=== Skipping playlists import ===")
            return

        print("=== Importing playlists ===")

        try:
            rs = self.piped.request("/user/playlists")
        except requests.exceptions.HTTPError as e:
            sys.stderr.write(f"===> [-] Could not fetch Piped playlists: {e}\n")
            sys.exit(1)

        piped_playlists = rs.json()
        print(f"=== Found {len(piped_playlists)} Piped playlists ===")

        try:
            rs = self.invidious.request("/auth/playlists")
        except requests.exceptions.HTTPError as e:
            sys.stderr.write(f"===> [-] Could not fetch Invidious playlists: {e}\n")
            sys.exit(1)

        invidious_playlists = {
            pl["title"]: {
                **pl,
                "videos": {item["videoId"]: item for item in pl.get("videos", [])},
            }
            for pl in rs.json()
        }

        print(f"=== Found {len(invidious_playlists)} Invidious playlists ===")

        for pl in piped_playlists:
            try:
                rs = self.piped.request(f'/playlists/{pl["id"]}')
            except requests.exceptions.HTTPError as e:
                sys.stderr.write(
                    f'===> [-] Could not fetch playlist: {pl["name"]}: {e}\n'
                )
                continue

            items = [
                re.sub(r"^/watch\?v=", "", item["url"])
                for item in rs.json().get("relatedStreams", [])
            ]

            print(f'Importing {len(items)} videos from playlist: {pl["name"]}')

            if pl["name"] in invidious_playlists:
                playlist_id = invidious_playlists[pl["name"]]["playlistId"]
                print(
                    f'[=] Playlist "{pl["name"]}" already exists on Invidious, skipping creation'
                )
            else:
                try:
                    rs = self.invidious.request(
                        "/auth/playlists",
                        method="POST",
                        json={
                            "title": pl["name"],
                            "privacy": "unlisted",
                        },
                    )
                except requests.exceptions.HTTPError as e:
                    sys.stderr.write(
                        f'===> [-] Could not create playlist on Invidious: {pl["name"]}: {e}\n'
                    )
                    continue

                playlist_id = rs.json().get("playlistId")
                if not playlist_id:
                    sys.stderr.write(
                        f'===> [-] Could not retrieve ID of the newly created Invidious playlist: {pl["name"]}\n'
                    )
                    continue

                print(
                    f'[+] Created playlist "{pl["name"]}" on Invidious, ID: {playlist_id}'
                )

            for item in items:
                if item in invidious_playlists.get(pl["name"], {}).get("videos", {}):
                    print(f'[=] Video "{item}" already exists on Invidious, skipping')
                    continue

                try:
                    rs = self.invidious.request(
                        f"/auth/playlists/{playlist_id}/videos",
                        method="POST",
                        json={"videoId": item},
                    )

                    print(f'[+] Added video "{item}" to playlist "{pl["name"]}"')
                except requests.exceptions.HTTPError as e:
                    sys.stderr.write(
                        f'===> [-] Could not add video {item} to playlist {pl["name"]}: {e}\n'
                    )
                    continue


class API(ABC):
    def __init__(self, session: Session, timeout: float = 10):
        self.session = session
        self.timeout = timeout

    @abstractmethod
    def request(self, url: str, method: str = "GET", **kwargs) -> requests.Response:
        pass


class InvidiousAPI(API):
    def request(self, url: str, method: str = "GET", **kwargs) -> requests.Response:
        rs = requests.request(
            method,
            urljoin(self.session.invidious_url, f'/api/v1/{url.lstrip("/")}'),
            headers={
                "Authorization": f"Bearer {self.session.invidious_token}",
            },
            **{
                "timeout": self.timeout,
                **kwargs,
            },
        )

        rs.raise_for_status()
        return rs


class PipedAPI(API):
    def request(self, url: str, method: str = "GET", **kwargs) -> requests.Response:
        rs = requests.request(
            method,
            urljoin(self.session.piped_api_url, url.lstrip("/")),
            headers={
                "Authorization": self.session.piped_token,
            },
            **{
                "timeout": self.timeout,
                **kwargs,
            },
        )

        rs.raise_for_status()
        return rs


def main():
    session = Session.from_cmdline(sys.argv[1:])
    session.migrate_subscriptions()
    session.migrate_playlists()


if __name__ == "__main__":
    main()