#!/usr/bin/env python3
# Author: Marcel-Brian Wilkowsky (mawigh)

import argparse;
import os;
import re;
import sys;
import json;
import time;
import threading;
import queue;
import shutil;
import requests;
import ventoyl;
from bs4 import BeautifulSoup;
from itertools import chain;
from difflib import SequenceMatcher;

def load (msg:str):
    characters = "/-\|";
    for char in characters:
        sys.stdout.write("\r"+msg+"... "+char);
        time.sleep(.1);
        sys.stdout.flush();

def hr_filesize (size):
    units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"];
    for unit in units:
        if size < 1024.0 or unit == "PiB":
            break;
        size /= 1024.0
    return f"{size:.0f} {unit}";

def download_file (mount_dir, iso_name, url):

    print("Downloading " + bchar.HEADER + str(iso_name) + bchar.ENDC + " from " + bchar.OKBLUE + str(url) + bchar.ENDC + " ...");
    download_request = requests.get(url + str(iso_name), stream=True, allow_redirects=True);
    start = time.process_time();
    download_length = download_request.headers.get("content-length");
    dl = 0;
    final_location = str(mount_dir) + "/" + str(iso_name);
    downloaded_file = open(final_location, "wb");

    for chunk in download_request.iter_content(1024):
        dl += len(chunk);
        downloaded_file.write(chunk);
        done = int(50 * dl / int(download_length));
        progress = "=" * done;
        half = " " * (50-done);
        timep = dl//(time.process_time() - start);
        timep = hr_filesize(timep/8);
        sys.stdout.write("\r"+bchar.OKCYAN+"["+bchar.ENDC + bchar.OKBLUE + str(progress) + str(half) +bchar.OKCYAN+"]"+bchar.ENDC+" "+str(timep));

    print("\nTime Elapsed: " + str(time.process_time() - start) + "s\n");
    downloaded_file.close();

class bchar:
    HEADER = '\033[95m'
    OKBLUE = '\033[94m'
    OKCYAN = '\033[96m'
    OKGREEN = '\x1b[32m'
    WARNING = '\033[93m'
    FAIL = '\033[91m'
    ENDC = '\033[0m'
    BOLD = '\033[1m'
    UNDERLINE = '\033[4m'

def main ():

    global etcdir;
    etcdir = os.path.abspath(os.path.join(os.path.dirname(ventoyl.__file__), "etc"));
    parser = argparse.ArgumentParser(prog="Ventoy Updater", description="Update ISO image found on a ventoy device");
    parser.add_argument("--device", help="explicit specifiy a Ventoy device (e.g. sda1)");
    parser.add_argument("-v", "--verbose", dest="v", action="store_true", help="enable verbose mode");
    parser.add_argument("-l", "--no-logo", dest="l", action="store_true", help="disable Ventoy ASCII logo");
    parser.add_argument("-V", "--version", action="version", version="%(prog)s 0.4");

    subparser = parser.add_subparsers(dest="command");

    sinstall = subparser.add_parser("install", help="install new ISO images");
    sinstall.add_argument("-m", "--no-umount", dest="m", action="store_true", help="DO NOT umount Ventoy device after install");
    sinstall.add_argument("-o", "--os", dest="short_name", help="short name of the OS you want to install (e.g. debian/manjaro)");
    sinstall.add_argument("-f", "--file", dest="file", help="copy specified ISO file to your Ventoy device (just a cp)");

    sremove = subparser.add_parser("remove", help="remove an installed ISO image from your Ventoy device");

    scheckurl = subparser.add_parser("check-url", description="If you leave the options empty, you will be asked interactively.", help="check an URL and return found ISO images");
    scheckurl.add_argument("--url", dest="url", help="the URL you want to check.", required=False);

    saddurl = subparser.add_parser("add-url", description="If you leave the options empty, you will be asked interactively.", help="add URL OS ISO sources to the JSON config file ("+ str(etcdir) +"/url_config.json)");
    saddurl.add_argument("--url", dest="addurl", help="the URL you want to add.", required=False);
    saddurl.add_argument("--short-name", dest="sname", help="the short name of your new entry. Like 'debian' or 'ubuntu'", required=False);
    saddurl.add_argument("--long-name", dest="lname", help="the long name of your added entry. Like 'Debian Buster' or 'Ubuntu Bionic Beaver'", required=False);

    scupdate = subparser.add_parser("check-update", help="Check for updates for your installed ISO images");

    supdate = subparser.add_parser("update", help="update all or specific ISO images");
    supdate.add_argument("--all", action="store_true", help="try to update all ISO images");

    slist = subparser.add_parser("list", help="list installed ISO images");
    sstats = subparser.add_parser("stats", help="print usage statistics");

    sconfig = subparser.add_parser("config", help="Modify your Ventoy device. Type config --help for more information");
    sconfigsubparser = sconfig.add_subparsers(dest="configcommand");

    sconfiginstall = sconfigsubparser.add_parser("install",  help="Install/Update latest Ventoy software on your device");
    sconfiginstall.add_argument("--gui", action="store_true", help="Opens the Ventoy GUI installer");
    sconfiginstall.add_argument("-f", "--force", action="store_true", help="Force the installation with the shell installer");

    sconfigmodify = sconfigsubparser.add_parser("modify",  help="Modify Ventoy attributes/plugins on your Ventoy device");
    sconfigmodify.add_argument("-f", "--force", action="store_true", help="Force the ...");

    smodifysubparser = sconfigmodify.add_subparsers(dest="modifycommand");

    smodifytheme = smodifysubparser.add_parser("theme", help="Changing a Ventoy theme");
    smodifytheme.add_argument("-l", "--list", action="store_true", help="list available themes (grub2)");

    slastlog = subparser.add_parser("lastlog", help="Output the last (10 lines) log information ("+str(etcdir)+"/ventoyl.log)");
    slastlog.add_argument("-l", "--lines", help="Specify how may lines should be printed to the output (default is 10 lines)");

    args = parser.parse_args();

    url_file = open(etcdir + "/url_config.json", "r");
    global json_data;
    json_data = json.load(url_file);

    global verbose_mode;
    verbose_mode = False;
    if args.v:
        verbose_mode = True;

    global ventoy;

    if not args.l:
        print(ventoyu_ascii);
    check_root();
    if args.device:
        try:
            ventoy = ventoyl.ventoyl(str(args.device), verbose_mode);
        except OSError:
            ventoy = None;
    else:
        ventoy = ventoyl.ventoyl(None, verbose_mode);

    #if ventoy:
    #    if not ventoy.ventoy_device:
    #        sys.exit( bchar.FAIL + "Error: No Ventoy device found. Plug in your Ventoy device or specify it explicity with --device");
    #else:
    #    sys.exit( bchar.FAIL + "Error: Are you sure your specified Ventoy device ("+str(args.device)+") exists? Please have a look.");
    if args.command == "stats":
        print_stats();
    if args.command == "install":
        umount = True;
        short_name = "";
        if args.m:
            umount = False;
        if args.short_name:
            short_name = args.short_name;
        install(umount, short_name);
    if args.command == "list":
        listfiles();
    if args.command == "update":
        update_iso();
    if args.command == "check-update":
        #load_wrapper("Checking for updates", check_update);
        check_update();
    if args.command == "remove":
        remove_iso();
    if args.command == "check-url":
        if args.url:
            checkurl(True, args.url);
        else:
            checkurl(False);
    if args.command == "add-url":
        print(args);
        if args.sname:
            s_name = args.sname;
        else:
            s_name = None;
        if args.lname:
            l_name = args.lname;
        else:
            l_name = None;
        if args.addurl:
            newurl = args.addurl;
        else:
            newurl = None;

        addurl(newurl, s_name, l_name);
    if args.command == "config":
        subaction = "";
        if args.configcommand == "modify":
            subaction = args.modifycommand;
            action = "modify";
        if args.configcommand == "install":
            if not args.device and not args.gui:
                sys.exit(bchar.FAIL + "Error: You have to explicity set your device with --device or start the GUI installer with --gui" + bchar.ENDC);
            action = "install";
            if args.gui:
                action = "opengui";
        if not args.configcommand:
            sconfig.print_help();
            sys.exit();
        config(action, args.force, subaction);
    if args.command == "lastlog":
        nlines = 10;
        if args.lines:
            nlines = args.lines;
        output_lastlog(nlines);
    if not args.command:
        print_stats();
        listfiles();

def check_root ():
    uid = os.getuid();
    if not uid == 0:
        sys.exit(bchar.FAIL + "Cannot access your Ventoy drive without root access.."+ bchar.ENDC +"\nTry again with"+ bchar.BOLD +" sudo " + " ".join(sys.argv)  + bchar.ENDC);
    else:
        return 0;

def print_stats ():

    if ventoy.ventoy_device:
        if ventoy.temp_dir:
            print(bchar.UNDERLINE + "Ventoy usage statistics:" + bchar.ENDC);
            print(bchar.BOLD + bchar.HEADER + "Total: " + bchar.ENDC + str(hr_filesize(shutil.disk_usage(ventoy.temp_dir).total)));
            print(bchar.BOLD + bchar.HEADER + "Used: " + bchar.ENDC + str(hr_filesize(shutil.disk_usage(ventoy.temp_dir).used)));
            print(bchar.BOLD + bchar.HEADER + "Free: " + bchar.ENDC + str(hr_filesize(shutil.disk_usage(ventoy.temp_dir).free)));

def load_wrapper (msg, func):
    process = threading.Thread(name="load_wrapper_process", target=func);
    process.start();
    while process.is_alive():
        load(msg);

def grequest (url):
    rval = requests.get(url).text;
    rqueue.put(rval);

def install(umount=True, short_name=""):
    
    json_iterator = 0;
    available_os = [];

    if short_name == "":    
        for os in json_data:
            print(bchar.BOLD + "[ID: "+ str(json_iterator) +"]"+bchar.ENDC+": " + bchar.OKGREEN + str(json_data[str(os)]["name"]) + bchar.ENDC)
            available_os.append(os);
            json_iterator += 1;
    
        select_install_os = input("Which OS do you want to install? [0-"+str(json_iterator-1)+"]: ");
    else:
        for os in json_data:
            available_os.append(os);
        try:
            select_install_os = str(available_os.index(str(short_name)));
        except ValueError:
            print(bchar.FAIL + "Error: Operating System " + bchar.BOLD + str(short_name) + bchar.ENDC + bchar.FAIL + " could not be found in your configuration file.\n"+bchar.HEADER+"Available are following:" + bchar.ENDC);
            print(", ".join([os for os in json_data]));
            sys.exit();

    try:
        try:
            url = json_data[available_os[int(select_install_os)]]["url"];

            #request = requests.get(url).text;
            global rqueue;
            rqueue = queue.Queue()
            request_process = threading.Thread(name="request_process", target=grequest, args=[str(url)]);
            request_process.start();
            while request_process.is_alive():
                load("Searching for installable versions");

            request_result = rqueue.get();

            sos_html_soup = BeautifulSoup(request_result, "html.parser");
            sos_iso_images = [ node.get("href") for node in sos_html_soup.find_all("a") if node.get("href").endswith(".iso")];
            sos_iso_images = list(dict.fromkeys(sos_iso_images))
            os_iterator = 0;
            print("\n");
            for image in sos_iso_images:
                image = image.split("/")[-1];
                print(bchar.BOLD + "[ID: "+ str(os_iterator) +"]"+bchar.ENDC+": " + bchar.OKGREEN + image + bchar.ENDC);
                os_iterator += 1;

            select_os_type = input("Which OS type do you want to install? [0-"+str(os_iterator-1)+"]: ");
            if select_os_type == "":
                if len(sos_iso_images) == 1:
                    select_os_type = "0";
    
            url_regex = r"^http[s]:\/\/.*\/";
            download_match = re.search(url_regex, sos_iso_images[int(select_os_type)]);
            if download_match:
                url = re.findall(url_regex, sos_iso_images[int(select_os_type)]);
                url = url[0];
                sos_iso_images[int(select_os_type)] = str(sos_iso_images[int(select_os_type)]).split("/")[-1];
            try:
                #download_url = url + sos_iso_images[int(select_os_type)];
                download_file(ventoy.temp_dir, sos_iso_images[int(select_os_type)], url);
                print(bchar.OKGREEN + "Done. You can now use "+str(sos_iso_images[int(select_os_type)])+"..." + bchar.ENDC);
                if umount:
                    umount_process = threading.Thread(name="umount_process", target=ventoy.umountVentoyDevice);
                    umount_process.start();
                    while umount_process.is_alive():
                        load("Unmounting " + str(ventoy.ventoy_device) + "... I may still have to sync.");
                    print("\n" + bchar.WARNING + "Ventoy device "+str(ventoy.ventoy_device)+ " umounted!" + bchar.ENDC);
            except IndexError or ValueError:
                sys.exit(bchar.FAIL + "Error: Select a valid operating system!" + bchar.ENDC);
        except KeyError or ValueError:
            sys.exit(bchar.FAIL + "Error: Interesting... Your JSON file looks broken. Cannot find parameter 'url' for os " + str(available_os[int(select_install_os)]) + bchar.ENDC);
    except IndexError or ValueError:
        sys.exit(bchar.FAIL + "Error: Select a valid operating system!" + bchar.ENDC);

def check_update (silent=False):

    _version_search_regex = r"([0-9]{1,4}\.[0-9]{1,3}\.iso)|(([0-9]{1,4}\.[0-9]{1,4})\.[0-9]{1,4})|([0-9]{1,4}\.)";

    iso_files = ventoy.getISOFiles();

    if len(iso_files) < 1:
        if not silent:
            sys.exit(bchar.FAIL + "Error: Could not find any ISO images on your Ventoy device!" + bchar.ENDC);

    check_for_updates = {};
    sum_iso_info = {};
    shortlist_update = {};

    for os in json_data:
        url = json_data[str(os)]["url"];
        request_result = requests.get(url).text;
    
        html_soup = BeautifulSoup(request_result, "html.parser");
        fiso_files = [ node.get("href") for node in html_soup.find_all("a") if node.get("href").endswith(".iso")];
        for fil in fiso_files:
            fil = fil.split("/")[-1];
            sum_iso_info.update({str(fil): str(url)});
    
    if not len(sum_iso_info) > 1:
        if not silent:
            sys.exit(bchar.FAIL + "Error: Could not retrieve any iso file from your config file!" + bchar.ENDC);

    for iso in iso_files:
        iso = iso.split("/")[-1];
        for siso, downloadurl in sum_iso_info.items():
            similarity = SequenceMatcher(None, siso, iso).ratio();
            if similarity == 1.0:
                if verbose_mode:
                    if not silent:
                        print( bchar.WARNING + "Verbose: " + bchar.OKGREEN + str(iso) + " is already the newest version!" + bchar.ENDC);
                break;
            if similarity > 0.89:
                shortlist_update.update({str(iso): [str(downloadurl), str(siso)]});
                if verbose_mode:
                    similarity = round(similarity*100, 1);
                    if not silent:
                        print(bchar.WARNING + "Verbose: Matching "+ iso +" and " + siso + " with "+bchar.BOLD+ str(similarity)+"%"+bchar.ENDC+bchar.WARNING+" similarity." + bchar.ENDC);

    for iso in shortlist_update:

        matched_similarity = shortlist_update[str(iso)][1];
        matched_url = shortlist_update[str(iso)][0];

        fregex = re.search(_version_search_regex, iso);
        sregex = re.search(_version_search_regex, matched_similarity);

        if fregex and sregex:
            frev = iso[:fregex.start()] + iso[fregex.end():];
            srev = matched_similarity[:sregex.start()] + matched_similarity[sregex.end():];
           
            similarity = SequenceMatcher(None, frev, srev).ratio();
            if similarity == 1.0:
                check_for_updates.update({str(iso): [shortlist_update[str(iso)][0], str(matched_similarity)]}); 

    if len(check_for_updates.items()) < 1:
        if not silent:
            if not verbose_mode:
                sys.exit(bchar.FAIL + "\nInfo: It seems that no ISO file on your Ventoy device is upgradable. You will get more information with " + bchar.BOLD + bchar.OKCYAN + "--verbose" + bchar.ENDC);
            else:
                sys.exit(bchar.FAIL + "\nInfo: It seems that no ISO file on your Ventoy device is upgradable." + bchar.ENDC);

    for i, x in check_for_updates.items():
        if not silent:
            print(bchar.HEADER + u"\u2192 " + str(i) + bchar.ENDC + bchar.OKGREEN + " seems to be upgradable!" + bchar.ENDC);

    if not silent:
        print(bchar.OKGREEN + "Type " + bchar.OKCYAN + bchar.BOLD + "sudo " + sys.argv[0].split("/")[-1] + " update " + bchar.ENDC + bchar.OKGREEN + "to update your ISO images."+ bchar.ENDC);

    return check_for_updates;

def remove_iso ():
    verbose_mode = True;
    listfiles();
    print("\nType 'ALL' to delete all ISO images.");
    which_iso = input("\nWhich ISO file do you want to remove? ");
    ventoy_files = ventoy.getISOFiles();
    if which_iso == "ALL":
        for iso in ventoy_files:
            ventoy.deleteISO(str(iso));
            print("Deleted ISO file "+ str(iso) +"!");

        return;

    try:
        returnval = ventoy.deleteISO(str(ventoy_files[int(which_iso)]));

    except ValueError:
        returnval = ventoy.deleteISO(str(which_iso));

    if returnval:
        print("Deleted ISO image.");
    else:
        print("Could not delete ISO image.");

def checkurl (explurl=False, url=None):

    checkurl = url;
    if not explurl:
        print("\nPlease type in the URL you want to check.");
        checkurl = input("URL: ");
    global rqueue;
    rqueue = queue.Queue()
    request_process = threading.Thread(name="request_process", target=grequest, args=[str(checkurl)]);
    request_process.start();
    while request_process.is_alive():
        load("Searching for installable versions");

    request_result = rqueue.get();

    sos_html_soup = BeautifulSoup(request_result, "html.parser");
    sos_iso_images = [ node.get("href") for node in sos_html_soup.find_all("a") if node.get("href").endswith(".iso")];
    os_iterator = 0;
    print("\n");
    for image in sos_iso_images:
        image = image.split("/")[-1];
        print(bchar.BOLD + "[ID: "+ str(os_iterator) +"]"+bchar.ENDC+": " + bchar.OKGREEN + image + bchar.ENDC);
        os_iterator += 1;

    if len(sos_iso_images) >= 1:
        print("\nI found "+ str(len(sos_iso_images)) +" files. You can add the URL by using following command:");
        print(bchar.HEADER + sys.argv[0] + " add-url --url " + checkurl + bchar.ENDC);
    else:
        print( bchar.WARNING + "Warning: I did not found any ISO images on " + bchar.HEADER + checkurl + bchar.ENDC + bchar.WARNING + ". Please try again with another URL." + bchar.ENDC);

def listfiles ():
    
    files_found = ventoy.getISOFiles();
    
    if files_found:
        print(bchar.UNDERLINE + "I found following ISO images on your Ventoy device:\n" + bchar.ENDC);
        for iso in files_found:
            size = os.stat(str(iso)).st_size;
            size = hr_filesize(size);
            sizel = "";
            if verbose_mode:
                sizel = bchar.HEADER + " ("+ str(size) +") " + bchar.ENDC;
            print(bchar.OKGREEN + u"\u2192 " + iso + bchar.ENDC + ""+sizel+" "+bchar.UNDERLINE+"["+str(files_found.index(iso))+"]" + bchar.ENDC);

        if verbose_mode:
            print("\nYou can use the index numbers to run an action (e.g. update or remove)");
    else:
        print(bchar.WARNING + bchar.UNDERLINE + "Warning:"+bchar.ENDC + bchar.WARNING + "\nNo ISO images on your Ventoy device ("+ventoy.ventoy_device+") found. Maybe its not a valid Ventoy device?" + bchar.ENDC);
        print(bchar.HEADER + "Use following command to install ISO images: " + bchar.ENDC + bchar.BOLD + "\nsudo ventoyu install" +bchar.ENDC);

def addurl (url, sname, lname):

    if isinstance(url, str):
        addurl = url;
    else:
        print("Please type in the URL you want to add to the JSON config file. You may want to check your URL with option check-url.");
        addurl = input("URL: ");
    
    if isinstance(sname, str):
        short_name = sname;
    else:
        print("Give your entry a short name like 'debian' or 'ubuntu'");
        short_name = input("Short name: ");

    if isinstance(lname, str):
        long_name = lname;
    else:
        print("Give your entry a long name like 'Debian Buster' or 'Manjaro KDE'");
        long_name = input("Long name: ");

    new_url_config = {
        str(short_name): {
            "name": str(long_name),
            "url": str(addurl)
        }
    };

    json_data.update(new_url_config);
    with open(etcdir + "/url_config.json", "w") as jsonconfig:
        json.dump(json_data, jsonconfig, indent=4);

def config (action, force=False, subaction=""):

    if action == "modify":
        ventoy.configureVentoyPlugin(subaction);
    if action == "install":
        print(bchar.HEADER + "Download latest Ventoy tarball to temp directory.. Trying to open the installer..." + bchar.ENDC);
        ventoy.installLatestVentoy(False, force);
    if action == "opengui":
        print(bchar.HEADER + "Download latest Ventoy tarball to temp directory.. Trying to open the GUI installer..." + bchar.ENDC);
        ventoy.installLatestVentoy(True);

def update_iso ():

    upgradable = check_update(True);
    
    if isinstance(upgradable, dict):
        if verbose_mode:
            print(bchar.WARNING + "Verbose: Found "+str(len(upgradable))+" upgradable images. Trying to update them..." + bchar.ENDC);

        for iso, isoinfo in upgradable.items():
            print(bchar.HEADER + "Updating "+ bchar.BOLD + bchar.OKBLUE + str(iso) + bchar.ENDC + bchar.HEADER + " with " + bchar.BOLD + bchar.OKBLUE +  str(isoinfo[1]) + bchar.ENDC + bchar.HEADER +  "..." + bchar.ENDC);

            if not ventoy.isVentoyMounted():
                ventoy.mountVentoyDevice();

            try:
                download_file(ventoy.temp_dir, isoinfo[1], isoinfo[0]);
            except:
                print(bchar.FAIL + "Error: Could not update " + bchar.BOLD + str(iso) + bchar.ENDC + bchar.FAIL + "!" + bchar.ENDC);
                continue;

            old_iso_path = ventoy.temp_dir + "/" + iso;
            if os.path.isfile(old_iso_path):
                rc = ventoy.deleteISO(str(iso));
                if rc == True:
                    if verbose_mode:
                        print(bchar.WARNING + "Info: Deleted the old ISO image " + str(iso) + bchar.ENDC);
                    print(bchar.OKGREEN + "Successfully updated " + bchar.BOLD + str(iso) + bchar.ENDC + bchar.OKGREEN + "!" + bchar.ENDC);
            else:
                print(bchar.FAIL + "Error: Could not delete the old ISO image " + str(iso) + bchar.ENDC);
                    

    if len(upgradable) < 1:
        sys.exit(bchar.FAIL + "Error: Could not found any upgradable images. Try run " + bchar.OKCYAN + bchar.BOLD + "sudo " + sys.argv[0].split("/")[-1] + " --verbose check-update"+bchar.ENDC + bchar.FAIL + " manually." + bchar.ENDC);

def output_lastlog (num_lines=10):

    log_file = ventoy.getVentoylLogFile();
    if log_file:
        rfile = open(log_file, "r");
        lines = rfile.readlines();
        for line in lines[-int(num_lines):]:
            print(line.rstrip("\n"));

ventoyu_ascii = bchar.OKBLUE + """
 _    __           __              __  __
| |  / /__  ____  / /_____  __  __/ / / /
| | / / _ \/ __ \/ __/ __ \/ / / / / / / 
| |/ /  __/ / / / /_/ /_/ / /_/ / /_/ /  
|___/\___/_/ /_/\__/\____/\__, /\____/   
                         /____/          
""" + bchar.ENDC;

if __name__ == "__main__":
    main();
