import os
import shutil
import hashlib
import json

def copy_folder_to_destination(src_path, dest_path, skip_dirs=None, skip_files=None):
    # Ensure the destination directory exists
    os.makedirs(os.path.dirname(dest_path), exist_ok=True)
    
    # Remove existing folder if present
    if os.path.exists(dest_path):
        shutil.rmtree(dest_path)
    
    # Create destination directory
    os.makedirs(dest_path)
    
    # Walk through source directory
    for root, dirs, files in os.walk(src_path):
        # Remove directories to skip from dirs list
        if skip_dirs:
            dirs[:] = [d for d in dirs if d not in skip_dirs]
            
        # Calculate relative path
        rel_path = os.path.relpath(root, src_path)
        dest_dir = os.path.join(dest_path, rel_path)
        
        # Create corresponding destination directory
        os.makedirs(dest_dir, exist_ok=True)
        
        # Copy files that aren't in skip_files
        for file in files:
            if skip_files and file in skip_files:
                continue
            src_file = os.path.join(root, file)
            dest_file = os.path.join(dest_dir, file)
            shutil.copy2(src_file, dest_file)

def calculate_md5(file_path):
    hash_md5 = hashlib.md5()
    with open(file_path, "rb") as f:
        for chunk in iter(lambda: f.read(4096), b""):
            hash_md5.update(chunk)
    return hash_md5.hexdigest()

def get_file_size(file_path):
    return os.path.getsize(file_path)

def update_json_file(json_array, folder_path):
    # Create new data array for files
    file_array = []
    
    # Walk through the copied folder and collect file details
    for root, _, files in os.walk(folder_path):
        for file in files:
            file_path = os.path.join(root, file)
            relative_path = os.path.relpath(file_path, folder_path)
            
            # Replace backslashes with forward slashes
            relative_path = relative_path.replace('\\', '/')
            
            file_entry = {
                "remote": os.path.join("data/", relative_path),
                "local": os.path.join("/", relative_path),
                "md5": calculate_md5(file_path),
                "size": get_file_size(file_path)
            }
            file_array.append(file_entry)
    
    # Replace the contents of the input json_array with new data
    json_array.clear()
    json_array.extend(file_array)

def update_files(src_path, dest_path, skip_dirs=None, skip_files=None):
    # Check if the source folder exists
    if not os.path.isdir(src_path) or not os.path.isdir(dest_path):
        print("Invalid folder path!")
        return
    
    # Copy all data contents
    copy_folder_to_destination(src_path, dest_path, skip_dirs, skip_files)
    print("Folder copied successfully.")

    
    
def main():
    
    here_path = os.path.dirname(os.path.abspath(__file__))
    project_path = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    
    # Path of the folder to copy (you can modify this)
    src_folder_name = "data"
    src_path = os.path.join(project_path, src_folder_name)
    #print(f"source path: {src_path}")
     
    # Path of the destination folder
    dest_folder_name = "firmware_update\\latest\\data"
    dest_path = os.path.join(project_path, dest_folder_name)
    #print(f"destination path: {dest_path}")
    
    # Path of the firmware binary file
    bin_name = ".pio\\build\\esp32s3dev\\firmware.bin"
    
    # Skip these directories
    skip_dirs = ["boards", "booths"]
    
    # Skip these files
    skip_files = ["wifi.json", "system.json", "luma-stiks.json", ]
    

    update_files = input("Do you want to update the files? (y/n): ")

    #   *********************** Copy Data Files ***********************
    if update_files.lower() == "y":
        update_files(src_path, dest_path, skip_dirs, skip_files)
    
        
    #   *********************** Copy Binary file ***********************
    update_firmware = input("Do you want to update the firmware? (y/n): ")
    if update_firmware.lower() == "y":
        # Copy firmware.bin to the destination
        bin_path = os.path.join(project_path, bin_name)
        shutil.copy(bin_path, here_path)
        #print(f"firmware path: {bin_path}")
        print("firmware.bin copied successfully.")
            
        
        
    #   *********************** Process update.json ***********************
    # Update the JSON file
    json_path = os.path.join(here_path, "update.json")
    print(f"json path: {json_path}")
    
    # Read existing JSON
    with open(json_path, "r") as f:
        try:
            json_doc = json.load(f)
        except json.JSONDecodeError:
            print("Invalid JSON file!")
            return
        
    # process the files array    
    #if update_files.lower() == "y":  
    json_files_array = json_doc["files"]
    update_json_file(json_files_array, dest_path)
        #print(f"Folder {os.path.basename(src_path)} processed successfully.")
        
    
    #   *********************** Process firmwware.bin in update.json ***********************
    # process the firmware 
    if update_firmware.lower() == "y":
        json_firmware = json_doc["firmware"]
        firmware_path = os.path.join(here_path, "firmware.bin")
        json_firmware["md5"] = calculate_md5(firmware_path)
        json_firmware["size"] = get_file_size(firmware_path)
        
    
    # Write updated JSON
    with open(json_path, "w") as f:
        json.dump(json_doc, f, indent=4)

    print("Update JSON files created successfully.")

if __name__ == "__main__":
    main()
