import os
import platform
import random
import re

import threading
from filelock import FileLock
import uuid
import zipfile

from time import sleep
import shutil
import psutil
import subprocess
import requests
from tqdm import tqdm
import docker

import logging

from desktop_env.providers.base import VMManager

logger = logging.getLogger("desktopenv.providers.vmware.VMwareVMManager")
logger.setLevel(logging.INFO)

MAX_RETRY_TIMES = 10
RETRY_INTERVAL = 5
UBUNTU_X86_URL = "docker-osworld-x86"

# Determine the platform and CPU architecture to decide the correct VM image to download
# if platform.system() == 'Darwin':  # macOS
#     # if os.uname().machine == 'arm64':  # Apple Silicon
#     URL = UBUNTU_ARM_URL
# # else:
# #     url = UBUNTU_X86_URL
# elif platform.machine().lower() in ['amd64', 'x86_64']:
#     URL = UBUNTU_X86_URL
# else:
#     raise Exception("Unsupported platform or architecture")

URL = UBUNTU_X86_URL

DOWNLOADED_FILE_NAME = URL.split('/')[-1]

if platform.system() == 'Windows':
    docker_path = r"C:\Program Files\Docker\Docker"
    os.environ["PATH"] += os.pathsep + docker_path

UBUNTU_X86_URL = "https://huggingface.co/datasets/xlangai/ubuntu_osworld/resolve/main/Ubuntu.qcow2"
VMS_DIR = "./vmware_vm_data"

def __download_vm(vms_dir: str):
    # Download the virtual machine image
    logger.info("Downloading the virtual machine image...")
    downloaded_size = 0

    URL = UBUNTU_X86_URL
    DOWNLOADED_FILE_NAME = URL.split('/')[-1]
    downloaded_file_name = DOWNLOADED_FILE_NAME

    while True:
        downloaded_file_path = os.path.join(vms_dir, downloaded_file_name)
        headers = {}
        if os.path.exists(downloaded_file_path):
            downloaded_size = os.path.getsize(downloaded_file_path)
            headers["Range"] = f"bytes={downloaded_size}-"

        with requests.get(URL, headers=headers, stream=True) as response:
            if response.status_code == 416:
                # This means the range was not satisfiable, possibly the file was fully downloaded
                logger.info("Fully downloaded or the file size changed.")
                break

            response.raise_for_status()
            total_size = int(response.headers.get('content-length', 0))

            with open(downloaded_file_path, "ab") as file, tqdm(
                    desc="Progress",
                    total=total_size,
                    unit='iB',
                    unit_scale=True,
                    unit_divisor=1024,
                    initial=downloaded_size,
                    ascii=True
            ) as progress_bar:
                try:
                    for data in response.iter_content(chunk_size=1024):
                        size = file.write(data)
                        progress_bar.update(size)
                except (requests.exceptions.RequestException, IOError) as e:
                    logger.error(f"Download error: {e}")
                    sleep(RETRY_INTERVAL)
                    logger.error("Retrying...")
                else:
                    logger.info("Download succeeds.")
                    break  # Download completed successfully

class DockerVMManager(VMManager):
    def __init__(self, registry_path=""):
        pass

    def get_vm_path(self, region):
        if not os.path.exists(os.path.join(VMS_DIR, DOWNLOADED_FILE_NAME)):
            __download_vm(VMS_DIR)
        return os.path.join(VMS_DIR, DOWNLOADED_FILE_NAME)