import logging
import os
import platform
import time
import zipfile
from typing import Literal

import requests
from tqdm import tqdm

logger = logging.getLogger()

MAX_RETRY_TIMES = 10
RETRY_INTERVAL = 5

UBUNTU_X86_URL = "https://huggingface.co/datasets/xlangai/ubuntu_osworld/resolve/main/Ubuntu.qcow2.zip"
WINDOWS_X86_URL = "https://huggingface.co/datasets/xlangai/windows_osworld/resolve/main/Windows-10-x64.qcow2.zip"


OSWORLD_DIR = os.environ.get("OSWORLD_VM_DIR", "").rstrip("/")
if not OSWORLD_DIR:
    OSWORLD_DIR = os.path.dirname(os.path.abspath(__file__))
VMS_DIR = os.path.join(OSWORLD_DIR, "docker_vm_data")

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


def _download_vm(vms_dir: str):
    global URL, DOWNLOADED_FILE_NAME
    # Download the virtual machine image
    logger.info("Downloading the virtual machine image...")
    downloaded_size = 0

    downloaded_file_name = DOWNLOADED_FILE_NAME

    os.makedirs(vms_dir, exist_ok=True)

    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}")
                    time.sleep(RETRY_INTERVAL)
                    logger.error("Retrying...")
                else:
                    logger.info("Download succeeds.")
                    break  # Download completed successfully

    if downloaded_file_name.endswith(".zip"):
        # Unzip the downloaded file
        logger.info("Unzipping the downloaded file...☕️")
        with zipfile.ZipFile(downloaded_file_path, "r") as zip_ref:
            zip_ref.extractall(vms_dir)
        logger.info("Files have been successfully extracted to the directory: " + str(vms_dir))


def get_osworld_vm_path(os_type: Literal["Ubuntu", "Windows"]):
    global URL, DOWNLOADED_FILE_NAME
    if os_type == "Ubuntu":
        URL = UBUNTU_X86_URL
    elif os_type == "Windows":
        URL = WINDOWS_X86_URL
    DOWNLOADED_FILE_NAME = URL.split("/")[-1]

    if DOWNLOADED_FILE_NAME.endswith(".zip"):
        vm_name = DOWNLOADED_FILE_NAME[:-4]
    else:
        vm_name = DOWNLOADED_FILE_NAME

    if not os.path.exists(os.path.join(VMS_DIR, vm_name)):
        _download_vm(VMS_DIR)
    return os.path.join(VMS_DIR, vm_name)
