#!/usr/bin/python3

import argparse
import hashlib
import os
import re
import shutil
import subprocess
import sys
import urllib.parse
import urllib.request
import zipfile
from pathlib import Path
from tqdm import tqdm

FREEDOS_DEFAULT_MIRROR = 'https://www.ibiblio.org/pub/micro/pc-stuff/freedos/files'

FREEDOS_BASE_URL = FREEDOS_DEFAULT_MIRROR + '/repositories'
FREEDOS11_URL = FREEDOS_BASE_URL + '/1.1'
FREEDOS12_URL = FREEDOS_BASE_URL + '/1.2'
FREEDOS13_URL = FREEDOS_BASE_URL + '/1.3'
FREEDOS_APPS = [ 'dn2' ]
FREEDOS_USERSPACE_TOOLS = [ 'assign', 'attrib', 'choice', 'comp', 'debug', 'display', 'edit', 'edlin', 'fc',
                    'find', 'format', 'htmlhelp', 'label', 'mem', 'mode', 'nansi', 'share', 'sort', 'swsubst', 'tree']
FREEDOS_ARCHIVES_KERNEL_COMMAND = [ 'kernel', 'command', 'freecom' ]
FREEDOS_ARCHIVES_EXTRA = ['defrag', 'deltree', 'diskcomp', 'diskcopy', 'exe2bin', 'more', 'move', 'replace',
                     'share', 'shsucdx', 'xcopy']
FREEDOS_UNIXLIKE_ARCHIVES = ['touch']
FREEDOS_UTIL_ARCHIVES = ['wcd']
FREEDOS_NET_ARCHIVES = ['curl', 'gopherus', 'links', 'ping', 'terminal']
FREEDOS_SOUND_ARCHIVES = ['dosmid', 'adplay', 'opencp']
FREEDOS_DEVEL_ARCHIVES = ['bwbasic']

TMP_DIR = os.path.join('/tmp', os.path.basename(sys.argv[0]) + '-' + os.environ['USER'] + '-' + str(os.getpid()))

def check_sum_pkg_info(destination_file):
    pkg_info_file = open(destination_file[:-3] + 'txt', 'r')
    unchecked = True
    for line in pkg_info_file.readlines():
        if line.startswith('SHA'):
            assert_sha256sum(destination_file, line[4:].strip())
            unchecked = False
            break
    if unchecked:
        print('Warning no SHA checksum in pkg-info. Could not verify checksum of ' + destination_file)

def check_sum(destination_file):
    if Path(destination_file[:-3] + 'txt').is_file():
        check_sum_pkg_info(destination_file)
    else:
        print('Warning could not verify checksum of ' + destination_file)

def download_file(source, destination_directory):
    destination_file = os.path.join(destination_directory, urllib.parse.unquote(source.split("/")[-1]))
    if Path(destination_file).is_file():
        if not destination_file.endswith('.txt'):
            check_sum(destination_file)
    else:
        print('Downloading ' + source + '...')
        with urllib.request.urlopen(urllib.parse.quote_plus(source, "\./_-:")) as response, open(destination_file, 'wb') as f:
            if hasattr(response, 'headers'):
                length = response.headers['Content-length']
            elif hasattr(response, 'getheader'):
                length = response.getheader('content-length')
            else:
                length = None
            if type(length) is str:
                with tqdm(total=int(length), unit='B', maxinterval=0.5) as pbar:
                    while True:
                        chunk = response.read(16384)
                        if not chunk:
                            break
                        f.write(chunk)
                        pbar.update(len(chunk))
            else:
                f.write(response.read())
    return destination_file

def download_pkg_info(freedos_url, packages, destination):
    for filename in packages:
        try:
            download_file(freedos_url + '/pkg-info/' + filename + '.txt', destination)
        except urllib.error.HTTPError:
            print('WARNING: The chosen version of FreeDOS does not contain ' + filename + '!')

def download_FREEDOS_USERSPACE_ARCHIVES(freedos_url, destination):
    print('Downloading FreeDOS userspace tools...')
    for filename in FREEDOS_APPS:
        try:
            download_file(freedos_url + '/apps/' + filename + '.zip', destination)
        except urllib.error.HTTPError:
            print('WARNING: The chosen version of FreeDOS does not contain ' + filename + '!')
    for filename in FREEDOS_USERSPACE_TOOLS:
        try:
            download_file(freedos_url + '/base/' + filename + '.zip', destination)
        except urllib.error.HTTPError:
            print('WARNING: The chosen version of FreeDOS does not contain ' + filename + '!')
    for filename in FREEDOS_UNIXLIKE_ARCHIVES:
        try:
            download_file(freedos_url + '/unix/' + filename + '.zip', destination)
        except urllib.error.HTTPError:
            # utilities that are now in 'unixline' were previously in 'util'
            FREEDOS_UTIL_ARCHIVES.append(filename)
    for filename in FREEDOS_UTIL_ARCHIVES:
        try:
            download_file(freedos_url + '/util/' + filename + '.zip', destination)
        except urllib.error.HTTPError:
            print('WARNING: The chosen version of FreeDOS does not contain ' + filename + '!')
    for filename in FREEDOS_NET_ARCHIVES:
        try:
            download_file(freedos_url + '/net/' + filename + '.zip', destination)
        except urllib.error.HTTPError:
            print('WARNING: The chosen version of FreeDOS does not contain ' + filename + '!')
    for filename in FREEDOS_SOUND_ARCHIVES:
        try:
            download_file(freedos_url + '/sound/' + filename + '.zip', destination)
        except urllib.error.HTTPError:
            print('WARNING: The chosen version of FreeDOS does not contain ' + filename + '!')
    for filename in FREEDOS_DEVEL_ARCHIVES:
        try:
            download_file(freedos_url + '/devel/' + filename + '.zip', destination)
        except urllib.error.HTTPError:
            print('WARNING: The chosen version of FreeDOS does not contain ' + filename + '!')

def download_FREEDOS_FULL_ARCHIVES(freedos_url, destination):
    print('Downloading FreeDOS kernel and FreeCOM...')
    for filename in FREEDOS_ARCHIVES_KERNEL_COMMAND:
        try:
            download_file(freedos_url + '/base/' + filename + '.zip', destination)
        except urllib.error.HTTPError:
            print('WARNING: The chosen version of FreeDOS does not contain ' + filename + '!')
    print('Downloading extra FreeDOS userspace tools')
    download_FREEDOS_USERSPACE_ARCHIVES(freedos_url, destination)

    print('Downloading done')

def calculate_sha256sum(filename):
    with open(filename,"rb") as f:
        return hashlib.sha256(f.read()).hexdigest();

def assert_sha256sum(file, sha256sum):
    val = calculate_sha256sum(file)
    if val == sha256sum:
        print('Verified checksum ' + str(file))
    else:
        print('The downloaded file ' + str(file) + ' could not be verified.')
        print('Actual   SHA256SUM: ' + val)
        print('Expected SHA256SUM: ' + sha256sum)
        print('Please remove the file and rerun the script.')
        sys.exit(1)

def download_files(imgurls, destination):
    destination_files = []
    for imgurl in imgurls:
        destination_files.append(download_file(imgurl, destination))
    return destination_files

def verify_noexistinginstall(destination):
    if os.listdir(destination) != []:
        print('There is a already an existing set of DOS installation files in ' + destination + '.')
        sys.exit(1)

def derive_url_list(imgurl):
    imgurls = []
    filename = urllib.parse.unquote(imgurl.split("/")[-1])
    if filename.casefold().find("disk") != -1:
        if filename.casefold().find("of") != -1:
            imgcnt = int(re.findall(r'(?i)disk\s?\d+\s?of\s(\d+)', filename)[0])
            for i in range(1, imgcnt+1):
                imgurls.append(imgurl.rsplit("/", 1)[0] + "/" + re.sub(r'(?i)(?P<one>disk\s?)\d+(?P<two>\s?of\s\d+)', r'\g<one>'+str(i) + r'\g<two>', filename))
        else: # if the filenames have just one number, assume the last disk from the set was provided
            imgcnt = int(re.findall(r'(?i)disk\s?(\d+)', filename)[0])
            for i in range(1, imgcnt+1):
                imgurls.append(imgurl.rsplit("/", 1)[0] + "/" + re.sub(r'(?i)(?P<one>disk\s?)\d+', r'\g<one>'+str(i), filename))
    else:
        imgurls.append(imgurl)
    return imgurls;

def download_dos(dos_flavour, destination):
    os.makedirs(destination, exist_ok=True)

    if dos_flavour in ['freedos13userspace', 'freedos13', 'freedosuserspace']:
        freedos_url = FREEDOS13_URL
    if dos_flavour in ['freedos12userspace', 'freedos12']:
        freedos_url = FREEDOS12_URL

    if dos_flavour in ['freedos12', 'freedos13']:
        download_pkg_info(freedos_url, FREEDOS_ARCHIVES_KERNEL_COMMAND + FREEDOS_UNIXLIKE_ARCHIVES + FREEDOS_UTIL_ARCHIVES + FREEDOS_NET_ARCHIVES, destination)
        download_FREEDOS_FULL_ARCHIVES(freedos_url, destination)
    if dos_flavour in ['freedos13userspace', 'freedos13', 'freedos12userspace', 'freedos12', 'freedosuserspace']:
        download_pkg_info(freedos_url, FREEDOS_USERSPACE_TOOLS, destination)
        download_FREEDOS_USERSPACE_ARCHIVES(freedos_url, destination)
        return
    if dos_flavour == 'freedos11':
        download_FREEDOS_FULL_ARCHIVES(FREEDOS11_URL, destination)
    if dos_flavour in ['freedos11userspace', 'freedos11']:
        download_FREEDOS_USERSPACE_ARCHIVES(FREEDOS11_URL, destination)
        return

parser = argparse.ArgumentParser(description="""Script to download either FreeDOS or a set of disk images.
Multiple versions of FreeDOS are preconfigured.
When a custom URL to a disk image is provided, based on known filename patterns, all related images are downloaded and extracted.
Either a name should contain \"Disk ? of x\" or \"diskx\" where 1 is the first disk and x is the last disk.
The destination directory will contain all files from all disk images.""")
parser.add_argument("-l", "--list", help="List available DOS variants", action="store_true")
parser.add_argument("-o", "--os", help="Download the specified DOS", type=str)
parser.add_argument("-d", "--destination", help="Specify/override the destination directory", type=str)
parser.add_argument("-s", "--sha256sum", help="Specify one or more sha256sum(s)", nargs='+', type=str)

args = parser.parse_args()
if args.list:
    print('freedos13 FreeDOS 1.3 (2022)')
    print('freedos13userspace FreeDOS 1.3 userspace (2022)')
    print('freedos12 FreeDOS 1.2 (2016)')
    print('freedos12userspace FreeDOS 1.2 userspace (2016)')
    print('freedos11 FreeDOS 1.1 (2011)')
    print('freedos11userspace FreeDOS 1.1 userspace (2011)')
    sys.exit(0)

if args.os:
    dos_to_download = args.os
    if args.destination:
        destination = args.destination
    else:
        destination = os.path.join(os.environ['HOME'], '.cache', 'dosemu', args.os.replace('userspace', ''))
    download_dos(dos_to_download, destination)
