#!/usr/bin/env python3
import requests
from sys import stderr, stdout, argv
from os import environ
from os.path import basename

host = environ.get("PLEX_HOST", "localhost")
token = environ.get("PLEX_TOKEN", "")

args = argv[1:]
DL = False
if args and args[0] == "-g":
    args = args[1:]
    DL = True

def prompt(text, prelude=None):
    global args
    if not args:
        if prelude:
            print(prelude, file=stderr)
        print(text, file=stderr, end="")
    else:
        if args[0] == "-s":
            args.pop(0)
            term1 = args.pop(0)
            term2 = None
            if "-" in term1:
                term1, term2 = term1.split("-",1)
            psplit = prelude.split("\n")
            match1 = psplit.index([m for m in psplit if term1.lower() in m.lower()][0])
            match2 = None
            if term2:
                match2 = psplit.index([m for m in psplit if term2.lower() in m.lower()][0])
            picked = f"{match1}{f'-{match2}' if match2 else ''}"
        else:
            picked = args.pop(0)
        return picked
    return input()

def get(endpoint):
    url = host + ("/" if endpoint[0] != "/" else "") + endpoint
    r = requests.get(url, params={"X-Plex-Token": token}, headers={"Accept": "application/json"})
    return r.json()

def main():
    libraries = [d for d in get("/library/sections")["MediaContainer"]["Directory"]]
    _ = []
    for l in sorted(libraries, key=lambda l: l["key"]):
        _.append(f"[{l['key']}] {l['title']}")
    choice = prompt("Pick library: ", prelude="\n".join(_))
    library = [l for l in libraries if l["key"] == choice][0]
    lib_list = [i for i in get(f"/library/sections/{library['key']}/all")["MediaContainer"]["Metadata"]]
    _ = []
    for n,i in enumerate(lib_list):
        _.append(f"[{n}] {i['title']}")
    picked = lib_list[int(prompt("Pick item: ", prelude="\n".join(_)))]

    if picked["type"] == "artist":
        tracks = [t for t in get(picked["key"].replace("children", "allLeaves"))["MediaContainer"]["Metadata"]]
        _ = []
        for n,t in enumerate(tracks):
            _.append(f"[{n}] {t['title']}")
        picked = tracks[int(prompt("Pick track: ", prelude="\n".join(_)))]
    elif picked["type"] == "show":
        episodes = [e for e in get(picked["key"].replace("children", "allLeaves"))["MediaContainer"]["Metadata"]]
        _ = []
        for n,e in enumerate(episodes):
            _.append(f"[{n}] S{e['parentIndex']:02d}E{e['index']:02d} {e['title']}")
        c = prompt("Pick episode: ", prelude="\n".join(_))
        if "-" in c:
            a,b = map(int,c.split("-"))
            picked = episodes[a:b+1]
        else:
            picked = [episodes[int(c)]]

    for item in picked:
        if len(item["Media"]) > 1:
            files = [f for f in item["Media"]]
            _ = []
            for n,f in enumerate(files):
                _.append(f"[{n}] {f['title']}")
            item = files[int(prompt("Pick file: ", prelude="\n".join(_)))]
        else:
            item = item["Media"][0]

        url = f"{host}{item['Part'][0]['key']}?X-Plex-Token={token}"
        if DL:
            fn = basename(item["Part"][0]["file"].replace("\\", "/"))
            print(f"{url}\n out={fn}")
        else:
            print(url)
    return 0

if __name__ == "__main__":
    from sys import exit
    exit(main())