#!/usr/bin/env python3

"""qbth.py - qbittorrent helper

Usage:
    qbth.py (HOSTNAME) (USERNAME) (PASSWORD)
    qbth.py -h

Examples:
    qbth.py "http://localhost:8080" "admin" "adminadmin"
    qbth.py "https://cat.seedhost.eu/lol/qbittorrent" "lol" "meow"

Options:
    -h, --help      show this help message and exit
"""

import json
import os
import subprocess
from shutil import which

import qbittorrentapi
import requests
from bs4 import BeautifulSoup
from docopt import docopt

args = docopt(__doc__)
conn_info = dict(
    host=args["HOSTNAME"],
    username=args["USERNAME"],
    password=args["PASSWORD"],
)
qbt_client = qbittorrentapi.Client(**conn_info)

try:
    qbt_client.auth_log_in()
except qbittorrentapi.LoginFailed as e:
    print(e)


def add_torrents(urls: list[str]):
    """
    Add torrents from their URLs.

    Params:
        urls: list of strings that are URLs.
    """
    with qbittorrentapi.Client(**conn_info) as qbt_client:
        for url in urls:
            response = requests.get(url)
            if response.status_code == 200:
                if qbt_client.torrents_add(url, category="distro") != "Ok.":
                    raise Exception("Failed to add torrent: " + os.path.basename(url))
                else:
                    print(f"Added {os.path.basename(url)}")
            else:
                print(f"{response.status_code}: {url}")


def add_torrents_from_html(webpage_url: str, torrent_substring: str):
    """
    Add torrents from an HTML web page.

    Params:
        webpage_url: a string that is the URL for the desired webpage.
        torrent_substring: a string that is a substring of the URLs in
            the webpage that you want to extract. It serves as a
            selector.
    """
    reqs = requests.get(webpage_url, timeout=60)
    soup = BeautifulSoup(reqs.text, "html.parser")
    with qbittorrentapi.Client(**conn_info) as qbt_client:
        for link in soup.find_all("a"):
            if torrent_substring in link.get("href"):
                url = f"{webpage_url}/{link.get('href')}"
                response = requests.get(url)
                if response.status_code == 200:
                    if qbt_client.torrents_add(url, category="distro") != "Ok.":
                        raise Exception(
                            "Failed to add torrent: " + os.path.basename(url)
                        )
                    else:
                        print(f"Added {os.path.basename(url)}")

                else:
                    print(f"{response.status_code}: {url}")


def remove_torrents(distro_substring: str):
    """
    Remove torrents by selecting a substring that corresponds to the
    distro's torrent file name. When the substring is found, the
    torrent is removed by passing the corresponding hash to the method.

    Params:
        distro_substring: a string that is a substring of the distro
            torrent's file name.
    """
    with qbittorrentapi.Client(**conn_info) as qbt_client:
        for torrent in qbt_client.torrents_info():
            if distro_substring in torrent.name:
                qbt_client.torrents_delete(
                    torrent_hashes=torrent.hash, delete_files=True
                )
                print(f"Removed {torrent.name}")


def add_almalinux(rel_ver: str):
    """
    Add AlmaLinux torrents from a list of URLs. These URLs are partially
    hardcoded for convenience and aren't expected to change frequently.

    Params:
        relver: the AlmaLinux release version.
    """
    urls = [
        f"https://almalinux-mirror.dal1.hivelocity.net/{rel_ver}/isos/aarch64/AlmaLinux-{rel_ver}-aarch64.torrent",
        f"https://almalinux-mirror.dal1.hivelocity.net/{rel_ver}/isos/ppc64le/AlmaLinux-{rel_ver}-ppc64le.torrent",
        f"https://almalinux-mirror.dal1.hivelocity.net/{rel_ver}/isos/s390x/AlmaLinux-{rel_ver}-s390x.torrent",
        f"https://almalinux-mirror.dal1.hivelocity.net/{rel_ver}/isos/x86_64/AlmaLinux-{rel_ver}-x86_64.torrent",
    ]

    add_torrents(urls)


def remove_almalinux(rel_ver: str):
    """
    Remove AlmaLinux torrents given their release version.

    Params:
        relver: the AlmaLinux release version.
    """
    remove_torrents(f"AlmaLinux-{rel_ver}")


def add_debian(rel_ver: str):
    """
    Add Debian torrents from a list of URLs.

    Params:
        relver: the Debian release version.
    """
    urls = [
        f"https://cdimage.debian.org/debian-cd/current/amd64/bt-dvd/debian-{rel_ver}-amd64-DVD-1.iso.torrent",
        f"https://cdimage.debian.org/debian-cd/current/arm64/bt-dvd/debian-{rel_ver}-arm64-DVD-1.iso.torrent",
        f"https://cdimage.debian.org/debian-cd/current/armel/bt-dvd/debian-{rel_ver}-armel-DVD-1.iso.torrent",
        f"https://cdimage.debian.org/debian-cd/current/armhf/bt-dvd/debian-{rel_ver}-armhf-DVD-1.iso.torrent",
        f"https://cdimage.debian.org/debian-cd/current/mips64el/bt-dvd/debian-{rel_ver}-mips64el-DVD-1.iso.torrent",
        f"https://cdimage.debian.org/debian-cd/current/mipsel/bt-dvd/debian-{rel_ver}-mipsel-DVD-1.iso.torrent",
        f"https://cdimage.debian.org/debian-cd/current/ppc64el/bt-dvd/debian-{rel_ver}-ppc64el-DVD-1.iso.torrent",
        f"https://cdimage.debian.org/debian-cd/current/s390x/bt-dvd/debian-{rel_ver}-s390x-DVD-1.iso.torrent",
    ]

    add_torrents(urls)


def remove_debian(rel_ver: str):
    """
    Remove Debian torrents given their release version.

    Params:
        relver: the Debian release version.
    """
    remove_torrents(f"debian-{rel_ver}")


def add_devuan(rel_ver: str):
    """
    Add Devuan torrents from a URL.

    Params:
        relver: the Devuan release version.
    """
    url = f"https://files.devuan.org/devuan_{rel_ver}.torrent"
    add_torrents([url])


def remove_devuan(rel_ver: str):
    """
    Remove Devuan torrents given their release version.

    Params:
        relver: the Devuan release version.
    """
    remove_torrents(f"devuan_{rel_ver}")


def add_fedora(rel_ver: str):
    """
    Add Fedora torrents from URLs extracted from a webpage.

    Params:
        relver: the Fedora release version.
    """
    webpage_url = "https://torrent.fedoraproject.org/torrents"
    torrent_substring = f"{rel_ver}.torrent"
    add_torrents_from_html(webpage_url, torrent_substring)


def remove_fedora(rel_ver: str):
    """
    Remove Fedora torrents given their release version.

    Params:
        relver: the Fedora release version.
    """
    with qbittorrentapi.Client(**conn_info) as qbt_client:
        for torrent in qbt_client.torrents_info():
            if torrent.name.startswith("Fedora") and torrent.name.endswith(rel_ver):
                qbt_client.torrents_delete(
                    torrent_hashes=torrent.hash, delete_files=True
                )
                print(f"Removed {torrent.name}")


def add_freebsd(rel_ver: str):
    """
    Add FreeBSD torrents via a text file on the web that contains their
    magnet links.

    Params:
        relver: the FreeBSD release version.
    """
    url = f"https://people.freebsd.org/~jmg/FreeBSD-{rel_ver}-R-magnet.txt"
    reqs = requests.get(url, timeout=60)
    data = reqs.text.split("\n")

    with qbittorrentapi.Client(**conn_info) as qbt_client:
        for line in data:
            if line.startswith("magnet:"):
                if qbt_client.torrents_add(line) != "Ok.":
                    raise Exception("Failed to add torrent: " + line.split("=")[2])

                print(f"Added {line.split('=')[2]}")


def remove_freebsd(rel_ver: str):
    """
    Remove FreeBSD torrents given their release version.

    Params:
        relver: the FreeBSD release version.
    """
    remove_torrents(f"FreeBSD-{rel_ver}")


def add_kali():
    """
    Add Kali Linux torrents from their URLs extracted from a webpage.
    This method does not accept any parameters. The latest Kali Linux
    version is automatically selected.

    Params: none
    """
    webpage_url = "https://kali.download/base-images/current"
    torrent_substring = ".torrent"
    add_torrents_from_html(webpage_url, torrent_substring)


def remove_kali():
    """
    Remove Kali Linux torrents. This method does not accept any parameters.
    All Kali Linux torrents in the qBittorrent instance will be removed.

    Params: none
    """
    remove_torrents("kali-linux")


def add_netbsd(rel_ver: str):
    """
    Add NetBSD torrents from their URLs extracted from a webpage.

    Params:
        relver: the NetBSD release version.
    """
    webpage_url = f"https://cdn.netbsd.org/pub/NetBSD/NetBSD-{rel_ver}/images/"
    torrent_substring = ".torrent"
    add_torrents_from_html(webpage_url, torrent_substring)


def remove_netbsd(rel_ver: str):
    """
    Remove NetBSD torrents given their release version.

    Params:
        relver: the NetBSD release version.
    """
    remove_torrents(f"NetBSD-{rel_ver}")


def add_nixos():
    """
    Add NixOS torrents from their GitHub release at
    https://github.com/AninMouse/NixOS-ISO-Torrents. This method does not
    accept any paramters. The latest NixOS torrent is automatically selected.

    Params: none
    """
    url = "https://api.github.com/repos/AnimMouse/NixOS-ISO-Torrents/releases/latest"
    reqs = requests.get(url, timeout=60)
    json_data = json.loads(reqs.text)

    with qbittorrentapi.Client(**conn_info) as qbt_client:
        for item in json_data["assets"]:
            response = requests.get(item["browser_download_url"])
            if response.status_code == 200:
                if (
                    qbt_client.torrents_add(
                        item["browser_download_url"], category="distro"
                    )
                    != "Ok."
                ):
                    raise Exception(
                        "Failed to add torrent: "
                        + os.path.basename(item["browser_download_url"])
                    )
                else:
                    print(f"Added {os.path.basename(item['browser_download_url'])}")
            else:
                print(f"{response.status_code}: {item['browser_download_url']}")


def remove_nixos():
    """
    Remove NixOS torrents. This method does not accept any parameters. All
    NixOS torrents in the qBittorrent instance will be removed.

    Params: none
    """
    remove_torrents("nixos")


def add_qubes(rel_ver: str):
    """
    Add QubesOS torrents from their URLs.

    Params:
        relver: the QubesOS release version.
    """
    url = f"https://mirrors.edge.kernel.org/qubes/iso/Qubes-R{rel_ver}-x86_64.torrent"

    response = requests.get(url)
    if response.status_code == 200:
        with qbittorrentapi.Client(**conn_info) as qbt_client:
            if qbt_client.torrents_add(url, category="distro") != "Ok.":
                raise Exception("Failed to add torrent: " + os.path.basename(url))
            else:
                print(f"Added {os.path.basename(url)}")
    else:
        print(f"{response.status_code}: {url}")


def remove_qubes(rel_ver: str):
    """
    Remove QubesOS torrents given their release version.

    Params:
        relver: the Qubes OS release version.
    """
    remove_torrents(f"Qubes-R{rel_ver}")


def add_rockylinux(rel_ver: str):
    """
    Add Rocky Linux torrents from their URLs.

    Params:
        relver: the Rocky Linux release version.
    """
    urls = [
        f"https://download.rockylinux.org/pub/rocky/{rel_ver}/isos/aarch64/Rocky-{rel_ver}-aarch64-dvd.torrent",
        f"https://download.rockylinux.org/pub/rocky/{rel_ver}/isos/ppc64le/Rocky-{rel_ver}-ppc64le-dvd.torrent",
        f"https://download.rockylinux.org/pub/rocky/{rel_ver}/isos/s390x/Rocky-{rel_ver}-s390x-dvd.torrent",
        f"https://download.rockylinux.org/pub/rocky/{rel_ver}/isos/x86_64/Rocky-{rel_ver}-x86_64-dvd.torrent",
    ]

    add_torrents(urls)


def remove_rockylinux(rel_ver: str):
    """
    Remove Rocky Linux torrents given their release version.

    Params:
        relver: the Rocky Linux release version.
    """
    remove_torrents(f"Rocky-{rel_ver}")


def add_tails(rel_ver: str):
    """
    Add Tails torrents from their URLs.

    Params:
        relver: the Tails release version.
    """
    urls = [
        f"https://tails.net/torrents/files/tails-amd64-{rel_ver}.img.torrent",
        f"https://tails.net/torrents/files/tails-amd64-{rel_ver}.iso.torrent",
    ]

    add_torrents(urls)


def remove_tails(rel_ver: str):
    """
    Remove Tails torrents given their release version.

    Params:
        relver: the Tails release version.
    """
    remove_torrents(f"tails-amd64-{rel_ver}")


if __name__ == "__main__":
    # Check if gum is installed.
    if which("gum") is None:
        exit("Please install gum first. https://github.com/charmbracelet/gum")

    # Run the gum program in a subprocess to allow easy selecting of distro
    # torrents.
    distro_selection = subprocess.run(
        [
            "gum",
            "choose",
            "--limit=1",
            "--header=Available torrents",
            "--height=13",
            "AlmaLinux",
            "Debian",
            "Devuan",
            "Fedora",
            "FreeBSD",
            "Kali Linux",
            "NetBSD",
            "NixOS",
            "Qubes",
            "Rocky Linux",
            "Tails",
        ],
        stdout=subprocess.PIPE,
        text=True,
        check=True,
    ).stdout.strip()

    # After the distro is selected and stored in the distro_selection variable,
    # choose an action to take on the selected distro.
    # Add: add the distro torrents to qBittorrent
    # Remove: remove the distro torrents from qBittorrent
    action_selection = subprocess.run(
        ["gum", "choose", "--limit=1", "--header='Choose:'", "Add", "Remove"],
        stdout=subprocess.PIPE,
        text=True,
        check=True,
    ).stdout.strip()

    # After the distro is selected and stored in the distro_selection variable,
    # and after the action is selected and store in the action_selection
    # variable, enter the release version of the selected distro to execute the
    # selected action on.
    relver = subprocess.run(
        [
            "gum",
            "input",
            f"--placeholder='Enter {distro_selection} release version'",
        ],
        stdout=subprocess.PIPE,
        text=True,
        check=True,
    ).stdout.strip()

    # Match the distro_selection to execute the action_selection on.
    match distro_selection:
        case "AlmaLinux":
            if action_selection == "Add":
                add_almalinux(relver)

            if action_selection == "Remove":
                remove_almalinux(relver)

        case "Debian":
            if action_selection == "Add":
                add_debian(relver)

            if action_selection == "Remove":
                remove_debian(relver)

        case "Devuan":
            if action_selection == "Add":
                add_devuan(relver)

            if action_selection == "Remove":
                remove_devuan(relver)

        case "Fedora":
            if action_selection == "Add":
                add_fedora(relver)

            if action_selection == "Remove":
                remove_fedora(relver)

        case "FreeBSD":
            if action_selection == "Add":
                add_freebsd(relver)

            if action_selection == "Remove":
                remove_freebsd(relver)

        case "Kali Linux":
            if action_selection == "Add":
                add_kali()

            if action_selection == "Remove":
                remove_kali()

        case "NetBSD":
            if action_selection == "Add":
                add_netbsd(relver)

            if action_selection == "Remove":
                remove_netbsd(relver)

        case "NixOS":
            if action_selection == "Add":
                add_nixos()

            if action_selection == "Remove":
                remove_nixos()

        case "Qubes":
            if action_selection == "Add":
                add_qubes(relver)

            if action_selection == "Remove":
                remove_qubes(relver)

        case "Rocky Linux":
            if action_selection == "Add":
                add_rockylinux(relver)

            if action_selection == "Remove":
                remove_rockylinux(relver)

        case "Tails":
            if action_selection == "Add":
                add_tails(relver)

            if action_selection == "Remove":
                remove_tails(relver)

        case _:
            print("Nothing to do.")