#!/usr/bin/env python3

from requests import Session
from os import utime
from os.path import exists, getmtime, getatime, getsize, abspath, join
import argparse
import re
from concurrent.futures import ThreadPoolExecutor

# args
parser = argparse.ArgumentParser(description='Download items from steam workshop.')
parser.add_argument('-i', '--items', metavar='ID', nargs='+', dest='item_ids', type=int, default=[], help='workshop item ids')
parser.add_argument('-c', '--collections', metavar='ID', nargs='+', dest='collection_ids', type=int, default=[], help='workshop collection ids')
parser.add_argument('-o', '--out', metavar='PATH', dest='out', type=abspath, default='.', help='output dir')
parser.add_argument('-f', '--filename', metavar='PATTERN', dest='filename', type=str, default='{title}-{id}.{ext}', help='output file basename pattern: defaults to \'{title}-{id}.{ext}\', also supports {filename}')
parser.add_argument('-p', '--parallel-downloads', metavar='NUMBER', dest='parallel_downloads', type=int, default=8, help='number of parallel downloads')
args = parser.parse_args()

if not args.item_ids and not args.collection_ids:
    parser.error('at least one of --items or --collections must be specified')

# init http session
session = Session()

def flash(string):
    print(f"\r\033[K{string}...", end="", flush=True)

def log(string):
    print(f"\r\033[K{string}")

# get items from collections
if args.collection_ids:
    flash('resolving collections')
    response = session.post(
        'http://api.steampowered.com/ISteamRemoteStorage/GetCollectionDetails/v1',
        data={
            'collectioncount' : len(args.collection_ids),
            **{
                f'publishedfileids[{i}]' : collection_id
                    for i, collection_id in enumerate(args.collection_ids)
            },
        },
    )
    response.raise_for_status()

    # extract item ids
    for collection in response.json()['response']['collectiondetails']:
        for child in collection.get('children', []):
            if child['filetype'] == 0:
                args.item_ids.append(int(child['publishedfileid']))

# get items
flash('resolving items')
response = session.post(
    'http://api.steampowered.com/ISteamRemoteStorage/GetPublishedFileDetails/v1',
    data={
        'itemcount' : len(args.item_ids),
        **{
            f'publishedfileids[{i}]' : plugin_id
                for i, plugin_id in enumerate(args.item_ids)
        },
    },
)
response.raise_for_status()

# create thread pool
executor = ThreadPoolExecutor(max_workers=args.parallel_downloads)

# download function
def download(item):
    # file found?
    if item['result'] != 1:
        raise ValueError(f"getting file '{item['publishedfileid']}' info failed: {item}")

    # get target path
    filename_without_ext, ext = item['filename'].split('/')[-1].rsplit('.', 1)
    target_path = join(args.out, f"{args.filename}".format(
        id=item['publishedfileid'],
        title=re.sub(r'[^a-zA-Z0-9]+', '_', item['title']).strip('_'),
        filename=filename_without_ext,
        ext=ext,
    ))

    # skip item?
    if (
        exists(target_path) and # exists
        int(item['time_updated']) == int(getmtime(target_path)) and # mtime matches
        int(item['file_size']) == int(getsize(target_path)) # filesize matches
    ):
        log(f'skipped {target_path}')
        return

    # download item
    response = session.get(item['file_url'], stream=True)
    response.raise_for_status()

    # one chunk at a time
    with open(target_path, 'wb') as file:
        written = 0
        for chunk in response.iter_content(chunk_size=1_048_576):
            if chunk:
                # print every 8MB
                written += len(chunk)
                if written // (8 * 1024 * 1024) > 0:
                    flash(f'downloading {target_path}')
                    written = 0

                file.write(chunk)

    # update modify time
    log(f'finished {target_path}')
    utime(target_path, (getatime(target_path), item['time_updated']))

# download items
for item in response.json()['response']['publishedfiledetails']:
    executor.submit(download, item)

# wait for all downloads to finish
executor.shutdown(wait=True)

# close session
session.close()