#####################################
##     PEinjector/installer        ##
#####################################

# This moudle is a part of PEinjector.
# Please comply with the LICENSE.

import os
import shutil
import json

############ Config ############
KEYWORD = "winpe"
PYTHON = "python"
DISK = "{disk}"
################################

# Linux please run: mount /dev/sdX /mnt/disk --mkdir
# DISK = "/mnt/disk"


def find_disk():  # This code from dashedgeless
    disk_list = "CDEFGHIJKABLMNOPQRSTUVWYZ"
    for i in disk_list:
        this_disk = i + ":"
        if os.path.exists(this_disk + "/ventoy/ventoy.json"):
            return this_disk
    raise OSError(
        "Cannot find the Ventoy disk!(try to run ventoyPlugson once)")


DISK = DISK.replace("{disk}", find_disk())

print("use disk: "+DISK)
with open(DISK+"/ventoy/ventoy.json", "r") as file:
    vtoycfg = json.load(file)
searchroot = ""
searchlevel = 3
if "control" in vtoycfg:
    for i in vtoycfg["control"]:
        if "VTOY_DEFAULT_SEARCH_ROOT" in i:
            searchroot = i["VTOY_DEFAULT_SEARCH_ROOT"]
        if "VTOY_MAX_SEARCH_LEVEL" in i:
            if i["VTOY_MAX_SEARCH_LEVEL"] != "max":
                searchlevel = int(i["VTOY_MAX_SEARCH_LEVEL"])
filelist = []


def dfs(deep, dir):
    if deep > searchlevel:
        return
    lists = os.listdir(DISK+dir)
    for i in lists:
        if os.path.isfile(DISK+dir+"/"+i):
            if ".iso" in i.lower() and KEYWORD in i.lower():
                filelist.append(dir+"/"+i)
        else:
            dfs(deep+1, dir+"/"+i)


dfs(0, searchroot)

if "injection" in vtoycfg:
    for i in vtoycfg["injection"]:
        if "image" in i:
            if i["image"] in filelist:
                filelist.remove(i["image"])
else:
    vtoycfg["injection"] = []
print("find iso:", filelist)
for i in filelist:
    vtoycfg["injection"].append({
        "image": i,
        "archive": "/PEinjector/PEinjector.7z"
    })
print("write config")
with open(DISK+"/ventoy/ventoy.json", "w") as file:
    file.write(json.dumps(vtoycfg))


print("copy file")
if not os.path.exists(DISK+"/PEinjector"):
    os.mkdir(DISK+"/PEinjector")
if os.path.exists(DISK+"/PEinjector/PEinjector.7z"):
    os.remove(DISK+"/PEinjector/PEinjector.7z")
shutil.copyfile("dist/PEinjector.7z", DISK+"/PEinjector/PEinjector.7z")
if not os.path.exists(DISK+"/PEinjector/disable.txt"):
    with open(DISK+"/PEinjector/disable.txt", "w") as file:
        pass

print("write version info")
shutil.copyfile("version", DISK+"/PEinjector/VERSION")

print("make dir")
if not os.path.exists(DISK+"/PEinjector/install"):
    os.mkdir(DISK+"/PEinjector/install")

print("install basic package")
if not os.path.exists(DISK+"/PEinjector/package"):
    os.mkdir(DISK+"/PEinjector/package")
for i in os.listdir("package"):
    if os.path.exists(DISK+"/PEinjector/package/"+i):
        shutil.rmtree(DISK+"/PEinjector/package/"+i)
    shutil.copytree("package/"+i, DISK+"/PEinjector/package/"+i)

print("done")