#!/usr/bin/env python3
import argparse
import os
from time import sleep

import toml
import transmission_rpc

# https://franklingu.github.io/programming/2016/03/01/creating-daemon-process-python-example-explanation/
def daemonize():
    try:
        pid = os.fork()
        if pid > 0:
            # exit first parent
            exit(0)
    except OSError:
        exit(1)

    # decouple from parent environment
    os.chdir("/")
    os.setsid()

    # do second fork
    try:
        pid = os.fork()
        if pid > 0:
            exit(0)
    except OSError:
        exit(1)

config = toml.load("/etc/nyaa-tools/config.toml")

def jelly_path(name: str, season: int, episode=None, offset=0) -> str:
    if episode is None:
        return "{base}/{name}/Season {season:02}".format(
            base=config["library"]["base_path"],
            name=name,
            season=season,
        )
    else:
        return "{base}/{name}/Season {season:02}/Episode S{season:02}E{episode:02}.mkv".format(
            base=config["library"]["base_path"],
            name=name,
            season=season,
            episode=episode + offset,
        )


parser = argparse.ArgumentParser(
    prog="nyaa-batch"
)
parser.add_argument('-t', '--title')
parser.add_argument('-s', '--season')
parser.add_argument('-l', '--link')
args = parser.parse_args()

try:
    season = int(args.season)
except TypeError:
    exit(1)


client = transmission_rpc.Client(host=config["transmission"]["host"])
torrent = client.add_torrent(args.link, files_wanted=None, paused=True)
torrent = client.get_torrent(torrent.id) # Don't ask
files = torrent.get_files()
files.sort()
wanted = []
for file in files:
    # Only get files in root
    if file.name.count('/') < 2:
        wanted.append((file.id, len(wanted) + 1))

client.change_torrent(torrent.id, files_wanted=map(lambda f: f[0], wanted))
client.start_torrent(torrent.id)
daemonize()

while len(wanted) > 0:
    torrent = client.get_torrent(torrent.id)
    get_files = torrent.get_files()
    for want in wanted:
        file = next(x for x in get_files if x.id == want[0])
        if file.size == file.completed:
            if not os.path.exists(jelly_path(args.title, season)):
                os.makedirs(jelly_path(args.title, season))
            os.link(
                config["transmission"]["download_path"] + "/" + file.name,
                jelly_path(args.title, season, want[1])
            )
            wanted.remove(want)
    sleep(5)