"""
╔════════════════════════════════════════╗
║                                        ║
║   Developed by Abdullah OZMANTAR       ║
║   GitHub : github.com/abdozmantar      ║
║                                        ║
╚════════════════════════════════════════╝
"""

import os
import subprocess
import venv
import importlib.metadata

requirements_file = "requirements.txt"
venv_dir = "venv"

if not os.path.exists(venv_dir):
    print(f"Creating virtual environment '{venv_dir}'...")
    venv.create(venv_dir, with_pip=True)
else:
    print(f"Virtual environment '{venv_dir}' already exists.")

if os.name == "nt":
    pip_executable = os.path.join(venv_dir, "Scripts", "pip.exe")
else:
    pip_executable = os.path.join(venv_dir, "bin", "pip")

def is_installed(package: str, version: str | None = None, strict: bool = True) -> bool:
    try:
        installed_version = importlib.metadata.version(package)
        if version:
            return installed_version == version if strict else installed_version.startswith(version)
        return True
    except importlib.metadata.PackageNotFoundError:
        return False
    except Exception as e:
        print(f"[Warning] Failed to check package '{package}': {e}")
        return False
    
def run_pip(*args):
    subprocess.run([pip_executable, 'install', '--no-warn-script-location', *args])


if os.path.exists(requirements_file):
    print(f"Installing packages from '{requirements_file}'...")
    run_pip("-r", requirements_file)
    print("Requirements installed successfully.")
else:
    print(f"'{requirements_file}' not found, skipping.")


if not is_installed("torch"):
    print("Installing PyTorch with CUDA support...")
    run_pip("torch", "torchaudio", "torchvision", "--index-url", "https://download.pytorch.org/whl/cu129")
else:
    import torch
    if not torch.cuda.is_available():
        print("PyTorch installed but CUDA not available. Installing CUDA-enabled version...")
        run_pip("torch", "torchaudio", "torchvision", "--index-url", "https://download.pytorch.org/whl/cu129")
    else:
        print("PyTorch with CUDA is already installed.")

print("\nSetup complete!")