import os
import platform
import sys
import ssl
import urllib.request
from pathlib import Path

def get_os_name():
 system = platform.system()
 if system == "Windows":
 return "Windows 64-bit" if platform.machine().endswith('64') else "Windows 32-bit"
 elif system == "Darwin":
 return "macOS"
 elif system == "Linux":
 return "Linux"
 elif system == "FreeBSD":
 return "FreeBSD"
 elif system.startswith("CYGWIN"):
 return "Cygwin"
 else:
 return "Other"

def download_file(url, output_filename):
 total_downloaded = 0
 total_truncated = 0 # only print download progress every 2MB to avoid spamming logs

 def report_progress(block_num, block_size, total_size):
 nonlocal total_downloaded
 nonlocal total_truncated
 total_downloaded += block_size
 if total_downloaded // (1024 * 1024) > total_truncated:
 total_truncated = total_downloaded // (1024 * 1024)
 if total_truncated % 2 == 0:
 print(f"\rDownloaded {total_downloaded // (1024 * 1024)} MB", end="")

 try:
 ssl._create_default_https_context = ssl._create_stdlib_context
 urllib.request.urlretrieve(url, output_filename, reporthook=report_progress)
 print(f"\nDownloaded {output_filename}")
 return True
 except Exception as e:
 print(f"\nFailed to download {output_filename}")
 print(f"Error: {str(e)}")
 sys.exit(1)

def check_os(os_name):
 print("\nChecking System")
 print("===============\n")
 print(f" Operating System : {os_name}")
 supported = {"macOS", "Linux"}
 if os_name not in supported:
 print("Unsupported operating system")
 sys.exit(1)

def download_dawn(os_name):
 print("\nDownload Dawn Library")
 print("=====================\n")

 outfile_map = {
 "macOS": "third_party/lib/libwebgpu_dawn.dylib",
 "Linux": "third_party/lib/libwebgpu_dawn.so",
 }
 url_map = {
 "macOS": "https://github.com/austinvhuang/dawn-artifacts/releases/download/0.2.0/libwebgpu_dawn.dylib",
 "Linux": "https://github.com/austinvhuang/dawn-artifacts/releases/download/0.2.0/libwebgpu_dawn.so",
 }

 outfile = outfile_map.get(os_name)
 url = url_map.get(os_name)

 if not outfile or not url:
 print(f"No download information for {os_name}")
 sys.exit(1)

 print(f" URL : {url}")
 print(f" Download File : {outfile}\n")
 print(" Downloading ...\n")

 if Path(outfile).exists():
 print(f" File {outfile} already exists, skipping.")
 sys.exit(0)

 Path(outfile).parent.mkdir(parents=True, exist_ok=True)
 download_file(url, outfile)

def setup_env(os_name):
 print("\nEnvironment Setup")
 print("=================\n")
 
 current_dir = os.getcwd()
 lib_dir = os.path.join(current_dir, "third_party", "lib")
 
 if os_name == "macOS":
 print(" Before running the program, run the following command or add it to your shell profile:")
 print(f" export DYLD_LIBRARY_PATH=$DYLD_LIBRARY_PATH:{lib_dir}")
 
 with open("source", "w") as f:
 f.write(f"export DYLD_LIBRARY_PATH=$DYLD_LIBRARY_PATH:{lib_dir}\n")
 if os_name == "Linux":
 print(" Before running the program, run the following command or add it to your shell profile:")
 print(f" export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:{lib_dir}")
 
 with open("source", "w") as f:
 f.write(f"export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:{lib_dir}\n")

def main():
 os_name = get_os_name()
 check_os(os_name)
 download_dawn(os_name)
 setup_env(os_name)
 print()

if __name__ == "__main__":
 main()