#! python3.11

#   MIT License - Copyright (c) 2023 - Captain FLAM
#
#   https://github.com/Captain-FLAM/KaraFan

import os, platform, shutil, subprocess, tempfile, zipfile

#************************************************************************
#****                         for DEVELOPERS                         ****
#************************************************************************

# This setup will Check if it was installed by "Git" --> No Auto-Update !

#************************************************************************

PYTHON		= "3.11"
CUDA		= "cu118"
ROCM		= "rocm5.6"  # Linux + AMD GPU

SYSTEM		= platform.system()
PIP			= ["py", "-" + PYTHON, "-m", "pip"] if SYSTEM == 'Windows' else ["python" + PYTHON, "-m", "pip"]

Repository  = 'https://github.com/Captain-FLAM/KaraFan'
Version_url = 'https://raw.githubusercontent.com/Captain-FLAM/KaraFan/master/App/__init__.py'
KaraFan		= 'https://codeload.github.com/Captain-FLAM/KaraFan/zip/refs/tags/v'
PyTorch		= "https://download.pytorch.org/whl/"
FFmpeg		= Repository + '/wiki/Data/FFmpeg_' + SYSTEM + '.zip'

# Get the current path of this script
Gdrive  = os.getcwd()
Project = os.path.join(Gdrive, "KaraFan")

Version = ""; Git_version = ""

def on_rm_error(func, path, exc_info):
	# path contains the path of the file that couldn't be removed
	# let's just assume that it's read-only and unlink it.
	if SYSTEM == 'Windows':
		subprocess.run(["attrib", "-r", path], text=True, capture_output=True, check=True)
	else:
		os.chmod(path, 0o777)  # Linux & Mac
	os.remove(path)

try:
	print('\nWelcome to the "KaraFan" Setup !\n')

	# Create missing folders
	user_folder = os.path.join(Gdrive, "KaraFan_user")
	os.makedirs(user_folder, exist_ok=True)
	os.makedirs(os.path.join(user_folder, "Models"), exist_ok=True)

	# Install Python minimal dependencies needed for this script !
	try:
		import requests
	except ImportError:
		try:
			subprocess.run(PIP + ["install", "requests"], text=True, capture_output=True, check=True)
			import requests
		except subprocess.CalledProcessError as e:
			raise Exception("Error during Install Python dependency : requests !", e.stderr, e.stdout)

	# Get the latest version on GitHub
	try:
		response = requests.get(Version_url)
		if response.status_code == requests.codes.ok:
			Git_version = response.text.splitlines()[0].replace("# Version", "").strip()
		else:
			raise Exception("Unable to check version on GitHub ! Maybe you're behind a firewall ?")
	except ValueError as e:
		raise Exception("Error when processing data !", e.stderr, e.stdout)
	except requests.exceptions.ConnectionError as e:
		raise Exception("Connection error while trying to fetch version !")

	# Auto-Magic update !

	Install = False
	if not os.path.exists(Project):
		Install = True  # First install !
	else:
		# Get local version
		local_file = os.path.join(Project, "App", "__init__.py")
		if not os.path.exists(local_file):
			raise Exception('Installation of "KaraFan" is corrupted ! Please delete JUST the folder "KaraFan" and try again.')
		
		with open(local_file, "r") as version_file:
			Version = version_file.readline().replace("# Version", "").strip()
		
		if Version != "" and Git_version != "":
			if Git_version > Version:
				print(f'A new version of "KaraFan" is available : {Git_version} !')
				
				# Check if it was installed by "Git" --> No Auto-Update !

				if os.path.exists(os.path.join(Project, ".git")):
					print("\nYou have to download the new version manually from :")
					print( Repository )
					print('... and compare to your code in "KaraFan" folder.\n')
				else:
					Install = True
			else:
				print('"KaraFan" is up to date.\n')

	if Install:
		print(f'Downloading the latest version of "KaraFan {Git_version}" from GitHub...')
		try:
			response = requests.get(KaraFan + Git_version, allow_redirects=True)
			if response.status_code == requests.codes.ok:
				# Create a temporary file
				temp = tempfile.NamedTemporaryFile(suffix='.zip', delete=False)
				temp.write(response.content)

				# Remove old files
				if os.path.exists(Project):
					shutil.rmtree(Project, onerror = on_rm_error)

				# Unzip the temporary file
				with zipfile.ZipFile(temp.name, 'r') as zip_ref:
					zip_ref.extractall(Gdrive)
					zip_ref.close()

				temp.close()
				os.remove(temp.name)

				# Rename the new folder
				shutil.move(os.path.join(Gdrive, "KaraFan-" + Git_version), Project)

				# Copy to the Parent directory
				if os.path.exists(os.path.join(Project, "KaraFan.pyw")):
					shutil.copyfile(os.path.join(Project, "KaraFan.pyw"), os.path.join(os.path.dirname(Project), "KaraFan.pyw"))
			else:
				raise Exception('Unable to download the latest version of "KaraFan" from GitHub !')
		except ValueError as e:
			raise Exception("Error when processing data !", e.stderr, e.stdout)
		except requests.exceptions.ConnectionError as e:
			raise Exception("Connection error while trying to connect !")
		
		print('"KaraFan" is installed.\n')

	os.chdir(Project)  # For pip install

	# Dependencies already installed ?
	print("Installing dependencies...")
	try:
		subprocess.run(PIP + ["install", "-r", "requirements.txt"], text=True, capture_output=True, check=True)
		subprocess.run(PIP + ["install", "-r", "requirements_PC.txt"], text=True, capture_output=True, check=True)
		
		print("Installation done.\n")

	except subprocess.CalledProcessError as e:
		raise Exception("Error during Install dependencies !", e.stderr, e.stdout)

	# Get FFmpeg from GitHub wiki
	ffmpeg_exe = os.path.join(user_folder, "ffmpeg") + (".exe" if SYSTEM == 'Windows' else "")

	if not os.path.exists(ffmpeg_exe):
		print('Downloading "FFmpeg"... -> ' + FFmpeg)
		try:
			response = requests.get(FFmpeg, allow_redirects=True)
			if response.status_code == requests.codes.ok:
				# Create a temporary file
				temp = tempfile.NamedTemporaryFile(suffix='.zip', delete=False)
				temp.write(response.content)

				# Unzip the temporary file
				with zipfile.ZipFile(temp.name, 'r') as zip_ref:
					zip_ref.extractall(user_folder)
					zip_ref.close()

				# Make it executable
				if platform.platform() == 'Linux':
					subprocess.run(["chmod", "777", ffmpeg_exe], text=True, capture_output=True, check=True)

				temp.close()
				os.remove(temp.name)
			else:
				raise Exception('Unable to download "FFmpeg" from GitHub wiki !')
		except ValueError as e:
			raise Exception("Error when processing data !", e.stderr, e.stdout)
		except requests.exceptions.ConnectionError as e:
			raise Exception("Connection error while trying to fetch FFmpeg !")
		
		print('"FFmpeg" downloaded.\n')

	# Install PyTorch CUDA
	
	Install = False
	try:
		import torch
		if SYSTEM != 'Darwin' and CUDA not in torch.__version__ and ROCM not in torch.__version__:
			subprocess.run(PIP + ["uninstall", "-y", "torch"], text=True, capture_output=True, check=True)
			Install = True
	except ImportError:
		Install = True  # Torch not installed

	if Install:
		print('Installing "PyTorch CUDA"... It will take a long time... Be patient !\n')
		try:
			# MPS acceleration is available on MacOS 12.3+
			if SYSTEM == 'Darwin':
				subprocess.run(PIP + ["install", "torch"], text=True, check=True)
			else:
				subprocess.run(PIP + ["install", "torch", "--index-url", PyTorch + CUDA], text=True, check=True)
			
		except subprocess.CalledProcessError as e:
			raise Exception('Error during Install "PyTorch CUDA" !', e.stderr, e.stdout)

	# Every step succeeded !
	print("\nEverything is Fine !\n")
	print('You can now run KaraFan by launching "KaraFan.py"\n')
	print('Simply double-click on it !\n')

except KeyboardInterrupt:
	print("\nInstallation Aborted !\n")

except Exception as e:
	print("")
	for line in e.args:  print(line)
	print("")

finally:
	# Wait a key to exit
	input("--> Press [Enter] to exit...\n")
	os._exit(0)