""" Main module of the installer subpackage for stouputils.

Provides functions for installing programs from local zip files or URLs.
It handles downloading, extracting, and setting up programs in a platform-agnostic way.

This module contains the core installation functions that are used by both the Windows
and Linux/macOS specific modules.
"""

# Imports
import os
import platform
import shutil
import tarfile
import zipfile
from collections.abc import Callable
from tempfile import TemporaryDirectory

import requests

from ..decorators import LogLevels, handle_error
from ..print.message import info, warning
from .common import *
from .linux import *
from .windows import *


# Helper functions
def extract_archive(
	extraction_path: str,
	temp_dir: str,
	extract_func: Callable[[str], None],
	get_file_list_func: Callable[[], list[str]]
) -> None:
	""" Helper function to extract archive files with consistent handling.

	Args:
		extraction_path     (str): Path where files should be extracted
		temp_dir            (str): Temporary directory for intermediate extraction
		extract_func        (Callable[[str], None]): Function to extract the archive
		get_file_list_func  (Callable[[], list[str]]): Function to get the list of files in the archive
	"""
	os.makedirs(extraction_path, exist_ok=True)

	# Check if the archive contains a single folder
	file_list: list[str] = get_file_list_func()
	root_dirs: set[str] = {item.split('/')[0] + '/' for item in file_list if '/' in item}

	# If all files are in a single root directory
	if len(root_dirs) == 1 and all(item.startswith(next(iter(root_dirs))) for item in file_list):
		# Extract to a temporary location first
		temp_extract_path: str = os.path.join(temp_dir, "extracted")
		os.makedirs(temp_extract_path, exist_ok=True)
		extract_func(temp_extract_path)

		# Move contents from the single folder to the final path
		single_folder_path: str = os.path.join(temp_extract_path, next(iter(root_dirs)).rstrip('/'))
		for item in os.listdir(single_folder_path):
			src_path: str = os.path.join(single_folder_path, item)
			dst_path: str = os.path.join(extraction_path, item)
			if os.path.isdir(src_path):
				shutil.copytree(src_path, dst_path, dirs_exist_ok=True)
			else:
				shutil.copy2(src_path, dst_path)
		debug(f"Extracted program contents from single folder to '{extraction_path}'")
	else:
		# Normal extraction if not a single folder
		extract_func(extraction_path)
		debug(f"Extracted program to '{extraction_path}'")


# Functions
def get_install_path(
	program_name: str,
	platform_str: str = platform.system(),
	ask_global: int = 0,
	add_path: bool = True,
	append_to_path: str = "",
	) -> str:
	""" Get the installation path for the program on the current platform.

	Args:
		program_name  (str):  The name of the program to install.
		platform_str  (str):  The platform to get the installation path for.
		ask_global    (int):  Whether to ask the user for a path, 0 = ask, 1 = install globally, 2 = install locally.
		add_path      (bool): Whether to add the program to the PATH environment variable.
		append_to_path (str):  String to append to the installation path when adding to PATH.
			(ex: "bin" if executables are in the bin folder)

	Returns:
		str: The installation path for the program.

	Examples:
		.. code-block:: python

			> # Get install path, ask user if they want global or local
			> path = get_install_path("waifu2x-ncnn-vulkan", ask_global=0)
			> # User will be prompted to choose between global and local installation

			> # Force local installation
			> path = get_install_path("my-tool", ask_global=2)
			> # Returns a local path like "/home/user/my-tool" or "C:/Users/user/my-tool"

			> # Force global installation
			> path = get_install_path("my-tool", ask_global=1)
			> # Returns a global path like "/usr/local/bin/my-tool" or "C:/Program Files/my-tool"
	"""
	platform_str = str(platform_str).lower()
	if platform_str == "windows":
		return get_install_path_windows(program_name, ask_global, add_path=add_path, append_to_path=append_to_path)
	elif platform_str == "linux":
		return get_install_path_linux(program_name, ask_global, add_path=add_path, append_to_path=append_to_path)
	else:
		warning(f"Unsupported platform for automatic install path: {platform_str}")
		return ""

def add_to_path(install_path: str, platform_str: str = platform.system()) -> bool:
	""" Add the program to the PATH environment variable.

	Args:
		install_path  (str):  The path to the program to add to the PATH environment variable.
		platform_str  (str):  The platform you are running on (ex: "Windows", "Linux", "Darwin", ...),
			we use this to determine the installation path if not provided.

	Returns:
		bool: True if add to PATH was successful, False otherwise.

	Examples:
		.. code-block:: python

			> # Add a directory to PATH on the current platform
			> success = add_to_path("/usr/local/bin/my-program")
			> print(success)
			True

			> # Add a Windows path (when running on Windows)
			> add_to_path("C:/Program Files/MyApp/bin", platform_str="Windows")
			True
	"""
	platform_str = str(platform_str).lower()
	if platform_str == "windows":
		return add_to_path_windows(install_path) is True
	elif platform_str == "linux":
		return add_to_path_linux(install_path) is True
	else:
		warning(f"Unsupported platform for automatic add to PATH: {platform_str}")
		return False

@handle_error(message="Failed during program installation", error_log=LogLevels.WARNING_TRACEBACK)
def install_program(
	input_path: str,
	install_path: str = "",
	platform_str: str = platform.system(),
	program_name: str = "",
	add_path: bool = True,
	append_to_path: str = "",
	) -> bool:
	""" Install a program to a specific path from a local zip file or URL.

	Args:
		input_path     (str):  Path to a zip file or a download URL.
		install_path   (str):  The directory to extract the program into, we ask user for a path if not provided.
		platform_str   (str):  The platform you are running on (ex: "Windows", "Linux", "Darwin", ...),
			we use this to determine the installation path if not provided.
		add_path       (bool): Whether to add the program to the PATH environment variable.
		program_name   (str):  Override the program name, we get it from the input path if not provided.
		append_to_path (str):  String to append to the installation path when adding to PATH.
			(ex: "bin" if executables are in the bin folder)

	Returns:
		bool: True if installation was successful, False otherwise.

	Examples:
		.. code-block:: python

			> # Install from a URL
			> success = install_program(
			>     "https://github.com/example/tool/releases/download/v1.0/tool.zip",
			>     install_path="/usr/local/bin/tool",
			>     program_name="tool",
			>     add_path=True
			> )
			> print(success)
			True

			> # Install from a local file
			> install_program(
			>     "/downloads/program.zip",
			>     install_path="C:/Program Files/MyProgram",
			>     add_path=False
			> )
			True

			> # Install with executables in a subdirectory
			> install_program(
			>     "https://example.com/ffmpeg.zip",
			>     program_name="ffmpeg",
			>     append_to_path="bin"  # Adds /path/to/ffmpeg/bin to PATH
			> )
			True
	"""
	# Get program name from input path if not provided
	# (ex: "https://example.com/program.zip" -> "program", "/var/www/program.exe" -> "program")
	if not program_name:
		program_name = os.path.splitext(os.path.basename(input_path))[0]

	# If no install path is provided, ask the user for one
	final_install_path: str = install_path
	if not install_path:
		final_install_path = get_install_path(program_name, platform_str, add_path=add_path, append_to_path=append_to_path)
		if not final_install_path:
			warning("Failed to get installation path, please provide a path to install the program to.")
			return False

	# Create a temporary directory
	with TemporaryDirectory() as temp_dir:
		temp_dir: str
		program_path: str = ""

		# Download the program if it's a URL
		if input_path.startswith("http"):
			info(f"Downloading program from '{input_path}'")
			response: requests.Response = requests.get(input_path)
			if response.status_code != 200:
				warning(f"Failed to download program from '{input_path}', reason: {response.reason}")
				return False

			# Save the program to a temporary directory
			temp_file: str = os.path.join(temp_dir, "program.zip")
			with open(temp_file, "wb") as f:
				f.write(response.content)
			program_path = temp_file
			debug(f"Downloaded program to '{program_path}'")
		else:
			program_path = input_path
			debug(f"Using local program path '{program_path}'")

		# Extract the program if it's a zip file
		if program_path.endswith(".zip"):
			with zipfile.ZipFile(program_path, "r") as zip_ref:
				extract_archive(
					final_install_path,
					temp_dir,
					lambda path: zip_ref.extractall(path),
					lambda: zip_ref.namelist()
				)

		# Extract the program if it's a tar.xz file
		elif program_path.endswith(".tar.xz"):
			with tarfile.open(program_path, "r:xz") as tar_ref:
				extract_archive(
					final_install_path,
					temp_dir,
					lambda path: tar_ref.extractall(path),
					lambda: tar_ref.getnames()
				)

		# Else, it's a directory so we just need to copy it
		elif os.path.isdir(program_path):
			shutil.copytree(program_path, final_install_path)

		else:
			warning(f"Failed to install program, input path is not a zip, tar.xz file or directory: '{program_path}'")
			return False

	# If add_path is True, and the installation path was provided, we add it to the PATH environment variable
	if add_path and install_path:
		if not add_to_path(os.path.join(final_install_path, append_to_path), platform_str):
			warning(
				f"Failed to add program to PATH, please add it manually to your PATH environment variable:\n"
				f"{final_install_path}"
			)
			return False

	# If we get here, the program was installed successfully
	return True

