import platform
import subprocess

import requests

# Define the path to the virtual machine
VM_PATH = r"Ubuntu\Ubuntu.vmx"  # change this to the path of your downloaded virtual machine


def download_and_unzip_vm():
    # Determine the platform and CPU architecture to decide the correct VM image to download
    if platform.system() == 'Darwin' and platform.machine() == 'arm64':  # macOS with Apple Silicon
        url = "https://huggingface.co/datasets/xlangai/ubuntu_arm/resolve/main/Ubuntu.zip"
    elif platform.system() in ['Linux',
                               'Windows'] and platform.machine() == 'x86_64':  # Linux or Windows with x86_64 CPU
        url = "https://huggingface.co/datasets/xlangai/ubuntu_x86/resolve/main/Ubuntu.zip"
    else:
        raise Exception("Unsupported platform or architecture")

    # Download the virtual machine image
    print("Downloading the virtual machine image...")
    subprocess.run(['wget', url], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
                   universal_newlines=True)

    # Unzip the downloaded file
    print("Unzipping the downloaded file...")
    subprocess.run(['unzip', 'Ubuntu.zip'], stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True,
                   universal_newlines=True)


# Execute the function to download and unzip the VM
download_and_unzip_vm()


# Determine the platform of the host machine and decide the parameter for vmrun
def get_vmrun_type():
    if platform.system() == 'Windows' or platform.system() == 'Linux':
        return '-T ws'
    elif platform.system() == 'Darwin':  # Darwin is the system name for macOS
        return '-T fusion'
    else:
        raise Exception("Unsupported operating system")


# Start the virtual machine
subprocess.run(f'vmrun {get_vmrun_type()} start "{VM_PATH}"', shell=True)
print("Virtual machine started...")

# Get the IP address of the virtual machine
get_vm_ip = subprocess.run(f'vmrun {get_vmrun_type()} getGuestIPAddress "{VM_PATH}" -wait', shell=True,
                           capture_output=True,
                           text=True)
print("Virtual machine IP address:", get_vm_ip.stdout.strip())

vm_ip = get_vm_ip.stdout.strip()


# Function used to check whether the virtual machine is ready
def download_screenshot(ip):
    url = f"http://{ip}:5000/screenshot"
    try:
        # max trey times 1, max timeout 1
        response = requests.get(url, timeout=(1, 1))
        if response.status_code == 200:
            return True
    except requests.exceptions.ConnectionError:
        pass
    return False


# Try downloading the screenshot until successful
while not download_screenshot(vm_ip):
    print("Check whether the virtual machine is ready...")

print("Virtual machine is ready. Start to make a snapshot on the virtual machine. It would take a while...")

# Create a snapshot of the virtual machine
subprocess.run(f'vmrun {get_vmrun_type()} snapshot "{VM_PATH}" "init_state"', shell=True)
print("Snapshot created.")