import os
import shutil
import re

def search_str(file_path, word):
    with open(file_path, 'r') as file:
        # read all content of a file
        content = file.read()
        # check if string present in a file
        return word in content

def replace_text(file_path, search_text, replace_text):
    with open(file_path, 'r+') as file:
        # Reading the content of the file
        # using the read() function and storing
        # them in a new variable
        data = file.read()
        # Searching and replacing the text
        # using the replace() function
        data = data.replace(search_text, replace_text)
        # Setting the position to the top
        # of the page to insert data
        file.seek(0)
        # Writing the replaced data in our
        # text file
        file.write(data)
        # Truncating the file size
        file.truncate()
        # Close file
        file.close()

def find_file(file_path, name):
    for dirpath, dirname, filename in os.walk(file_path):
        for fn in filename:
            if name in fn:
                return os.path.join(dirpath, fn)

def write_file(file_path, name, content):
    with open(os.path.join(file_path, name), 'w') as file:
        # Writing the replaced data in our
        # text file
        file.write(content)
        # Truncating the file size
        file.truncate()
        # Close file
        file.close()

# Data dict
data = {
    "js": {
        "basePath": "/usr/share/truenas/webui",
        "expressions": [
            { "searchRegex": "if([...t].length>=(null===(i=this.availableGpus)",
              "replaceRegex": "if([...t].length>(null===(i=this.availableGpus)"},
            { "searchRegex": "if(t(a).length>=(null===(o=e.availableGpus)",
              "replaceRegex": "if(t(a).length>(null===(o=e.availableGpus)"},
            { "searchRegex": "{name:\"gpus\"});if(i.length&&i.length>=o.options.length)",
              "replaceRegex": "{name:\"gpus\"});if(i.length&&i.length>o.options.length)"},
            { "searchRegex": "{name:\"gpus\"});if(r.length&&r.length>=c.options.length)",
              "replaceRegex": "{name:\"gpus\"});if(r.length&&r.length>c.options.length)"},
        ]
    },
    "python": {
        "basePath": "/usr/lib/python3/dist-packages/middlewared/plugins/system_advanced",
        "baseDeletePath": "/usr/lib/python3/dist-packages/middlewared/plugins/system_advanced/__pycache__/",
        "expressions": [
            { "searchRegex": "            if len(available - provided) < 1:\n                verrors.add(\n                    f'{schema}.isolated_gpu_pci_ids',\n                    'A minimum of 1 GPU is required for the host to ensure it functions as desired.'\n                )",
              "replaceRegex": "#            if len(available - provided) < 1:\n#                verrors.add(\n#                    f'{schema}.isolated_gpu_pci_ids',\n#                    'A minimum of 1 GPU is required for the host to ensure it functions as desired.'\n#                )" }
        ]
    }
}

# Find JS files containing relevant lines
for file in os.listdir(data['js']['basePath']):
    if file.endswith(".js"):
        for entry in data['js']['expressions']:
            if search_str(os.path.join(data['js']['basePath'], file), entry['searchRegex']):
                entry.update({"path": os.path.join(data['js']['basePath'], file)})
                #print(entry['searchRegex'] + " found in: " + os.path.join(data['js']['basePath'], file))

# Backup JS files
for entry in data['js']['expressions']:
    if 'path' in entry:
        if not os.path.exists(entry['path'] + ".orig"):
            shutil.copy2(entry['path'], entry['path'] + ".orig")
        else:
            print(entry['path'] + " already backed up")

# Replace JS file contents
for entry in data['js']['expressions']:
    if 'path' in entry:
        replace_text(entry['path'], entry['searchRegex'], entry['replaceRegex'])

# Find Python files containing relevant lines
for file in os.listdir(data['python']['basePath']):
    if file.endswith(".py"):
        for entry in data['python']['expressions']:
            if search_str(os.path.join(data['python']['basePath'], file), entry['searchRegex']):
                entry.update({"path": os.path.join(data['python']['basePath'], file)})
                #print(entry['searchRegex'] + " found in: " + os.path.join(data['python']['basePath'], file))

# Backup Python files
for entry in data['python']['expressions']:
    if 'path' in entry:
        if not os.path.exists(entry['path'] + ".orig"):
            shutil.copy2(entry['path'], entry['path'] + ".orig")
        else:
            print(entry['path'] + " already backed up") 

# Replace Python file contents
for entry in data['python']['expressions']:
    if 'path' in entry:
        replace_text(entry['path'], entry['searchRegex'], entry['replaceRegex'])

# Remove Python cache files
cache_file = find_file(data['python']['baseDeletePath'], "config")
if cache_file and os.path.exists(cache_file):
    data['python']['cache_file'] = cache_file
    os.remove(cache_file)

# Generate restore shell script
restore_script = "#!/bin/bash\n\n"
for entry in data['js']['expressions']:
    if 'path' in entry:
        restore_script += "cp " + entry['path'] + ".orig " + entry['path'] + "\n"
for entry in data['python']['expressions']:
    if 'path' in entry:
        restore_script += "cp " + entry['path'] + ".orig " + entry['path'] + "\n"
for entry in data['python']:
    if 'cache_file' in entry and entry['cache_file'] is not None:
        restore_script += "rm " + entry['cache_file'] + "\n"
        
print("Writing restore script ...")

write_file(os.getcwd(), "restore_single_gpu_passthrough_changes.sh", restore_script)

print("Restarting services ...")

# Restart middlewared and nginx
os.system("service middlewared restart")
os.system("service nginx restart")

print("Patch successful!")