import subprocess
import sys
import asyncio
import os

async def run_remote_command(remote_host, command):
    ssh_command = ['ssh', f'root@{remote_host}', command]
    process = await asyncio.create_subprocess_exec(
        *ssh_command,
        stdout=asyncio.subprocess.PIPE,
        stderr=asyncio.subprocess.PIPE
    )
    stdout, stderr = await process.communicate()

    stdout_decoded = stdout.decode()
    stderr_decoded = stderr.decode()

    if process.returncode == 0:
        return stdout_decoded + stderr_decoded
    else:
        return f"Failed to run the command: {command}, error: {stderr_decoded}"

async def check_GX_type(remote_host):
    command = "cat /etc/venus/machine"
    return await run_remote_command(remote_host, command)

async def resize(remote_host):
    command = "sh /opt/victronenergy/swupdate-scripts/resize2fs.sh"
    return await run_remote_command(remote_host, command)

async def update_cerbo_firmware(remote_host):
    command = "sh /opt/victronenergy/swupdate-scripts/check-updates.sh -swu http://updates.victronenergy.com/feeds/venus/release/images/einstein/venus-swu-einstein-20240523125018-v3.32.swu"
    return await run_remote_command(remote_host, command)

async def set_official_update_feed(remote_host):
    command = "dbus -y com.victronenergy.settings /Settings/System/ReleaseType SetValue %0"
    return await run_remote_command(remote_host, command)

async def enable_large_image(remote_host):
    command = "dbus -y com.victronenergy.settings /Settings/System/ImageType SetValue %1"
    return await run_remote_command(remote_host, command)

async def check_large_image_or_not(remote_host):
    command = "dbus -y com.victronenergy.settings /Settings/System/ImageType GetValue"
    return await run_remote_command(remote_host, command)

async def update_firmware(remote_host):
    command = "sh /opt/victronenergy/swupdate-scripts/check-updates.sh -update"
    return await run_remote_command(remote_host, command)

async def enable_NodeRed(remote_host):
    command = "dbus -y com.victronenergy.platform /Services/NodeRed/Mode SetValue %1"
    return await run_remote_command(remote_host, command)

async def disable_NodeRed(remote_host):
    command = "dbus -y com.victronenergy.platform /Services/NodeRed/Mode SetValue %0"
    return await run_remote_command(remote_host, command)

async def check_NodeRed_enabled_or_not(remote_host):
    command = "dbus -y com.victronenergy.platform /Services/NodeRed/Mode GetValue"
    return await run_remote_command(remote_host, command)

async def download_node_red_dashboard(remote_host):
    change_dir_command = "cd /data/home/nodered/.node-red"
    install_command = "npm install --no-audit --no-update-notifier --no-fund --save --save-prefix=~ --production --engine-strict node-red-dashboard@latest"
    command = f"{change_dir_command} && {install_command}"
    return await run_remote_command(remote_host, command)

async def disable_BMS_Controlling_MPPT(remote_host):
    command = "dbus -y com.victronenergy.solarcharger /Settings/BmsPresent SetValue %0"
    return await run_remote_command(remote_host, command)

async def check_mkVersion(remote_host):
    command = "dbus -y com.victronenergy.vebus.ttyS4 /Interfaces/Mk2/Version GetValue"
    return await run_remote_command(remote_host, command)

async def check_allow_mkVersion_update_or_not(remote_host):
    command = "dbus -y com.victronenergy.settings /Settings/Vebus/AllowMk3Fw212Update GetValue"
    return await run_remote_command(remote_host, command)

async def update_mkVersion(remote_host):
    command = "dbus -y com.victronenergy.settings /Settings/Vebus/AllowMk3Fw212Update SetValue %1"
    return await run_remote_command(remote_host, command)

async def import_pika(remote_host):
    change_dir_command = "cd /data/innovenergy/pika-0.13.1/"
    install_command = "python3 setup.py install --user"
    command = f"{change_dir_command} && {install_command}"
    return await run_remote_command(remote_host, command)

async def reboot(remote_host):
    command = "reboot"
    return await run_remote_command(remote_host, command)

async def upload_files(remote_host,which_file):
    if which_file == 1:
        file_location_mappings = {
            "rc.local": "/data/",
            "dbus-fzsonick-48tl": "/data/",
            "aggregator": "/data/",
            "service": "/data/innovenergy/openvpn/",
            "openvpn": "/data/innovenergy/openvpn/",
            "pika-0.13.1": "/data/innovenergy/"
        }
    else: 
        file_location_mappings = {
            "flows.json": "/data/home/nodered/.node-red/",
            "settings-user.js": "/data/home/nodered/.node-red/",
        }

    cerbo_release_files_folder = os.path.join(os.getcwd(), "CerboReleaseFiles")
    if not os.path.exists(cerbo_release_files_folder):
        return "CerboReleaseFiles folder does not exist."

    try:
        tasks = []
        for file_name, remote_location in file_location_mappings.items():
            file_path = os.path.join(cerbo_release_files_folder, file_name)
            if not os.path.exists(file_path):
                raise FileNotFoundError(f"File {file_name} not found in {cerbo_release_files_folder}.")

            command = [
                "rsync",
                "-r",
                file_path,
                f"root@{remote_host}:{remote_location}"
            ]

            tasks.append(command)

        # Execute rsync commands asynchronously
        for task in tasks:
            subprocess.run(task, check=True)

        return "All files uploaded successfully."

    except FileNotFoundError as e:
        return str(e)

    except subprocess.CalledProcessError as e:
        return f"Error occurred while uploading files: {e}"

    except Exception as e:
        return f"An error occurred while uploading files: {str(e)}"
    
async def check_connection(remote_host):
    # Check if the remote host is reachable
    ping_command = ['ping', '-c', '1', remote_host]
    result = subprocess.run(ping_command, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    return result.returncode == 0

async def main(remote_host):
    
    ##################################################### Part 1 ################################################################
    ##### 1. upload VPN and battery files ######
    print("Upload VPN and battery files!")
    if(await upload_files(remote_host,1)!="All files uploaded successfully."):
        sys.exit("Failed to upload files!")
    else:
        print(await upload_files(remote_host,1))

    ##### 2. update firmware with normal image #####
    gx_type = await check_GX_type(remote_host)
    if gx_type == "einstein\n":
        print("Update Cerbo GX firmware now!")
        print(await update_cerbo_firmware(remote_host))
    else:
        sys.exit("It's not Cerbo GX!")
    ##################################################### Part 1 ################################################################

    # Wait for remote computer to come back online
    while not await check_connection(remote_host):
        print("Waiting for remote computer to come back online...")
        await asyncio.sleep(10)  # Adjust timeout as needed

    ##################################################### Part 2 ################################################################
    #### 3. udpate to MK3 #####
    if (await check_mkVersion(remote_host) == "value = 1170212\n" and await check_allow_mkVersion_update_or_not(remote_host) == "0\n"):
        print("Update MK3!")
        print(await update_mkVersion(remote_host))
    else:
        print("No need to update to MK3!")

    #### 4. import pika ####
    print("Import pika!")
    print(await import_pika(remote_host))

    ##### 5. update firmware with large image #####
    print("Set update feed to official release now!")
    print(await set_official_update_feed(remote_host))
    print("Enable large image now!")
    print(await enable_large_image(remote_host))
    image_type = await check_large_image_or_not(remote_host)
    if image_type == "1\n":
        print("Update firmware with large image!")
        print(await update_firmware(remote_host))
    else:
        sys.exit("Failed to enable large image!")
    ##################################################### Part 2 ################################################################

    # Wait for remote computer to come back online
    while not await check_connection(remote_host):
        print("Waiting for remote computer to come back online...")
        await asyncio.sleep(10)  # Adjust timeout as needed

    ##################################################### Part 3 ################################################################
    #### 6. resize /dev/root #####
    print("Resize /dev/root now!")
    print(await resize(remote_host))

    #### 7. enable Node Red #####
    print("Enable Node Red now!")
    print(await enable_NodeRed(remote_host))
    if(await check_NodeRed_enabled_or_not(remote_host) == "value = 1\n"):
        ##### 8. download Node Red Dashboard #####
        print("Download Node Red Dashboard now!")
        print(await download_node_red_dashboard(remote_host))
    else:
        sys.exit("Failed to enable Node Red!")

    ##### 9. upload files related to Node Red #####
    print("Upload files related to Node Red now!")
    if(await upload_files(remote_host,2)!="All files uploaded successfully."):
        sys.exit("Failed to upload files!")

    #### 10. restart Node Red to load and deploy flows #####
    print("Disable Node Red!")
    if(await disable_NodeRed(remote_host) == "retval = 0\n"):
        print("Re-enable Node Red!")
        print(await enable_NodeRed(remote_host))
        if(await check_NodeRed_enabled_or_not(remote_host) == "value = 1\n"):
            print("Node Red is set now!")
        else:
            sys.exit("Failed to re-enable Node Red!")
    ##################################################### Part 3 ################################################################
   
if __name__ == "__main__":
    remote_host = sys.argv[1]
    asyncio.run(main(remote_host))