akibabot/akibabot/main.py

import sys
import os
import json
from pathlib import Path
from dataclasses import dataclass
import logging
import asyncio

import toml
import akibapass_downloader.login
import akibapass_downloader.episode

from akibabot.xmpp import send_notification


@dataclass
class Show:
    # List of episode numbers that are locally available.
    local_episodes: list[int]

    # The name of the show.
    name: str

    # The release year of the show.
    year: int

    # The season number, if specified
    season: int | None

    # The base URL of the show.
    url: str


def build_show_dir_name(name: str, year: int) -> str:
    """Build the Jellyfin-compatible directory name of the show."""
    return f"{name} ({year})"


def build_episodes_path(
    name: str, year: int, season: int | None, storage_path: Path
) -> Path:
    show_path = storage_path / build_show_dir_name(name, year)

    if season is not None:
        return show_path / f"Season {season}"
    return show_path


def get_episode_numbers(
    show_name: str,
    show_year: int,
    show_season: int | None,
    storage_path: Path,
) -> list[int]:
    """Computes the episode numbers that are locally available.

    @show_name: The name of the show
    @show_year: The release year of the show.
    @storage_path: Path to where the show directories are.
    """
    episodes_path = build_episodes_path(show_name, show_year, show_season, storage_path)
    if not episodes_path.exists():
        return []

    return [
        int(item.split(".")[0].split(" ")[1])
        for item in os.listdir(episodes_path)
        if (episodes_path / item).is_file()
    ]


def setup_logging():
    """Sets up the logging."""
    logging.basicConfig(level=logging.INFO)


def main():
    # Setup logging
    setup_logging()

    # Load config
    config_file = sys.argv[1]
    config = toml.load(config_file)
    if (secrets_file := config["general"].get("secrets_file", None)) is not None:
        logging.info("Processing secrets config file...")
        secrets = toml.load(secrets_file)
        config.update(secrets)

    # Sanity checks
    xmpp_from_jid = config["xmpp"]["jid"]
    xmpp_to_jid = config["xmpp"]["to"]
    xmpp_password = config["xmpp"]["password"]

    # First, filter shows that are not done.
    storage_path = Path(config["general"]["path"])
    shows: list[Show] = []
    for show in config["shows"]:
        logging.info("Processing %s", show["name"])
        max_episodes = show["max_episodes"]
        local_episodes = get_episode_numbers(
            show["name"],
            show["year"],
            show.get("season", None),
            storage_path,
        )
        logging.debug("=> %d/%d", max(local_episodes), max_episodes)
        if len(set(local_episodes)) != max_episodes:
            shows.append(
                Show(
                    local_episodes=list(set(local_episodes)),
                    name=show["name"],
                    year=show["year"],
                    url=show["url"],
                    season=show.get("season", None),
                ),
            )

    # Exit early when there is no show left to handle.
    if not shows:
        print("Nothing to do.")
        sys.exit(0)

    # Load cookies, if available
    cookies = {}
    cookies_path = Path(config["general"]["cookie_path"])
    if cookies_path.exists():
        with open(cookies_path, "r", encoding="utf8") as f:
            cookies = json.load(f)

    # Request a new session, if required
    if (not cookies) or not akibapass_downloader.login.is_logged_in(cookies):
        logging.info("Requesting new cookies...")
        cookies = akibapass_downloader.login.login(
            config["user"]["email"],
            config["user"]["password"],
        )
        with open(cookies_path, "w", encoding="utf8") as f:
            f.write(json.dumps(cookies))
        logging.info("Done")

    # Iterate over the shows to process and download the new episodes
    for show in shows:
        logging.debug("Processing %s", show.name)
        episodes_path = build_episodes_path(
            show.name,
            show.year,
            show.season,
            storage_path,
        )
        if not episodes_path.exists():
            logging.info(
                "Episodes directory of %s does not exist. Creating...", show.name
            )
            episodes_path.mkdir(parents=True)

        episodes_remote = akibapass_downloader.episode.list_episodes(
            show.url,
        )
        episodes_to_download = set(
            [episode.episode_nr for episode in episodes_remote]
        ) - set(show.local_episodes)

        for episode in episodes_remote:
            if episode.episode_nr in episodes_to_download:
                logging.info("Downloading %s (%d)", episode.name, episode.episode_nr)
                downloads = episode.get_downloads(
                    cookies=cookies,
                    filter_quality=akibapass_downloader.episode.Quality.UHD_1440P,
                )
                if not downloads:
                    logging.warning(
                        "Failed to find UHD quality for episode %d", episode.episode_nr
                    )
                    continue

                download = downloads[0]
                download.download(
                    cookies=cookies,
                    destination_dir=episodes_path,
                    filename=f"Episode {episode.episode_nr}.mp4",
                )
                logging.info("Download of episode %s done", episode.episode_nr)
                # Send a message to notify us.
                asyncio.get_event_loop().run_until_complete(
                    send_notification(
                        xmpp_from_jid,
                        xmpp_password,
                        xmpp_to_jid,
                        show.name,
                        episode.episode_nr,
                    ),
                )