import launch
import importlib
import sys
from packaging.version import Version
from packaging.requirements import Requirement

def is_installed(pip_package):
    """
    Check if a package is installed and meets version requirements specified in pip-style format.

    Args:
        pip_package (str): Package name in pip-style format (e.g., "numpy>=1.22.0").

    Returns:
        bool: True if the package is installed and meets the version requirement, False otherwise.
    """
    try:
        # Parse the pip-style package name and version constraints
        requirement = Requirement(pip_package)
        package_name = requirement.name
        specifier = requirement.specifier  # e.g., >=1.22.0

        # Check if the package is installed
        dist = importlib.metadata.distribution(package_name)
        installed_version = Version(dist.version)

        # Check version constraints
        if specifier.contains(installed_version):
            return True
        else:
            print(f"Installed version of {package_name} ({installed_version}) does not satisfy the requirement ({specifier}).")
            return False
    except importlib.metadata.PackageNotFoundError:
        print(f"Package {pip_package} is not installed.")
        return False


# Define requirements directly
requirements = [
    "einops>=0.8.1",
    "fvcore",
    "easydict",
    "matplotlib>=3.3.0",
    "yacs",
    "scikit-image",
    "wandb",
    "torch>=2.4.1",
    "torchvision>=0.11.0",
    "timm>=1.0.20",
    "numpy>=1.19.0",
    "Pillow>=8.0.0",
    "scikit-learn>=1.2.0",
    "scipy>=1.8.0",
    "tqdm>=4.60.0",
    "safetensors>=0.4.0"
]

# Add triton-windows only on Windows
if sys.platform == 'win32':
    requirements.append("triton-windows")

# Add webui packages
webui_packages = [
    "gradio>=3.41,<6",
    "fastapi",
    "uvicorn",
    "python-multipart"
]
requirements.extend(webui_packages)

for req in requirements:
    if not is_installed(req):
        launch.run_pip(f"install {req}", f"comfyui-kaloscope requirement: {req}")