import logging
import os
import platform
import time
import docker
import psutil
import requests
from filelock import FileLock
from pathlib import Path

from desktop_env.providers.base import Provider

logger = logging.getLogger("desktopenv.providers.docker.DockerProvider")
logger.setLevel(logging.INFO)

WAIT_TIME = 3
RETRY_INTERVAL = 1
LOCK_TIMEOUT = 10


class PortAllocationError(Exception):
    pass


class DockerProvider(Provider):
    def __init__(self, region: str):
        self.client = docker.from_env()
        self.server_port = None
        self.vnc_port = None
        self.chromium_port = None
        self.environment = {"DISK_SIZE": "32G", "RAM_SIZE": "4G", "CPU_CORES": "4"}  # Modify if needed

        temp_dir = Path(os.getenv('TEMP' if platform.system() == 'Windows' else '/tmp'))
        self.lock_file = temp_dir / "docker_port_allocation.lck"
        self.lock_file.parent.mkdir(parents=True, exist_ok=True)

    def _get_available_port(self, port: int, lock_file: Path = None):
        if lock_file is None:
            lock_file = self.lock_file
        lock = FileLock(str(lock_file), timeout=LOCK_TIMEOUT)
        with lock:
            while port < 65354:
                if port not in [conn.laddr.port for conn in psutil.net_connections()]:
                    return port
                port += 1

    def start_emulator(self, path_to_vm: str, headless: bool, os_type: str):
        self.vnc_port = self._get_available_port(8006)
        self.server_port = self._get_available_port(5000)
        # self.remote_debugging_port = self._get_available_port(1337)
        self.chromium_port = self._get_available_port(9222)
        logger.info(f"Occupying ports: {self.vnc_port}, {self.server_port}, {self.chromium_port}")
        self.container = self.client.containers.run("happysixd/osworld-docker", environment=self.environment,
                                                    cap_add=["NET_ADMIN"], devices=["/dev/kvm"], volumes={
                os.path.abspath(path_to_vm): {"bind": "/System.qcow2", "mode": "ro"}},
                                                    ports={8006: self.vnc_port, 5000: self.server_port,
                                                           9222: self.chromium_port}, detach=True)

        def download_screenshot(ip, port):
            url = f"http://{ip}:{port}/screenshot"
            try:
                # max trey times 1, max timeout 1
                response = requests.get(url, timeout=(10, 10))
                if response.status_code == 200:
                    return True
            except Exception as e:
                time.sleep(RETRY_INTERVAL)
            return False

        # Try downloading the screenshot until successful
        while not download_screenshot("localhost", self.server_port):
            logger.info("Check whether the virtual machine is ready...")

    def get_ip_address(self, path_to_vm: str) -> str:
        return f"localhost:{self.server_port}:{self.chromium_port}:{self.vnc_port}"

    def save_state(self, path_to_vm: str, snapshot_name: str):
        raise NotImplementedError("Not available for Docker.")

    def revert_to_snapshot(self, path_to_vm: str, snapshot_name: str):
        pass

    def stop_emulator(self, path_to_vm: str):
        logger.info("Stopping VM...")
        self.container.stop()
        self.container.remove()
        time.sleep(WAIT_TIME)