#!/usr/bin/env bash
# shellcheck disable=SC2086,SC2155,SC2046,SC1091,SC2317

# =============================================================================
# DeepGEMM Installer Script
# =============================================================================
#
# This script installs DeepGEMM in a robust, idempotent way. It:
# - Works when executed directly or from another bash script
# - Uses repo-relative paths by default (no reliance on CWD)
# - Optionally clones DeepGEMM (HTTPS by default) if not present
# - Prefers the local `apps/DeepGEMM` directory if it exists
# - Calls upstream install.sh and verifies Python package import
# - Provides clear logs and non-zero exit on failure
#
# Usage:
#   ./install_deepgemm.sh [options]
#
# Options:
#   -p, --project_path PATH   Project root (default: auto-detected from this script)
#   -r, --repo URL            DeepGEMM repo URL (default: https://github.com/deepseek-ai/DeepGEMM.git)
#   -b, --branch BRANCH       Branch/Tag to checkout when cloning (default: main)
#   -d, --dir PATH            Install source directory (default preference: apps/DeepGEMM or $PROJECT_PATH/DeepGEMM)
#   -f, --force-clone         Force cloning even if local apps/DeepGEMM exists
#   -h, --help                Show help
#
# Examples:
#   ./install_deepgemm.sh
#   ./install_deepgemm.sh -p /path/to/mt_dao
#   ./install_deepgemm.sh -r https://github.com/deepseek-ai/DeepGEMM.git -b v0.1.0
# =============================================================================

set -Eeuo pipefail

readonly SCRIPT_NAME="install_deepgemm.sh"

# Pretty logging
log_info() { echo "$SCRIPT_NAME: $1"; }
log_warn() { echo "$SCRIPT_NAME: WARNING: $1" >&2; }
log_error() { echo "$SCRIPT_NAME: ERROR: $1" >&2; }

# Trap for unexpected failures
err_trap() {
	local rc=$?
	log_error "Failed at line $LINENO (exit=$rc)"
	exit "$rc"
}

trap 'err_trap' ERR

# Defaults
readonly DEFAULT_REPO_URL=${DEEPGEMM_REPO:-"https://github.com/deepseek-ai/DeepGEMM.git"}
readonly DEFAULT_BRANCH=${DEEPGEMM_BRANCH:-"main"}

# Compute repo root (two levels up from this script: scripts/system_setup/.. -> project root)
_THIS_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
readonly DEFAULT_PROJECT_PATH="$(cd "$_THIS_DIR/../.." && pwd)"

PROJECT_PATH="$DEFAULT_PROJECT_PATH"
DEEPGEMM_REPO="$DEFAULT_REPO_URL"
DEEPGEMM_BRANCH="$DEFAULT_BRANCH"
INSTALL_DIR="" # If empty, will resolve later to apps/DeepGEMM or $PROJECT_PATH/DeepGEMM
FORCE_CLONE=false

show_usage() {
	cat <<EOF
Usage: $SCRIPT_NAME [OPTIONS]

Install DeepGEMM from local apps/DeepGEMM or by cloning the upstream repo.

OPTIONS:
	-p, --project_path PATH   Project root (default: $DEFAULT_PROJECT_PATH)
	-r, --repo URL            DeepGEMM repo URL (default: $DEFAULT_REPO_URL)
	-b, --branch BRANCH       Branch/Tag to checkout when cloning (default: $DEFAULT_BRANCH)
	-d, --dir PATH            Install source directory (default: prefer apps/DeepGEMM, else clone to $PROJECT_PATH/DeepGEMM)
	-f, --force-clone         Force cloning even if local apps/DeepGEMM exists
	-h, --help                Show this help message

EXAMPLES:
	$SCRIPT_NAME
	$SCRIPT_NAME -p /custom/mt_dao
	$SCRIPT_NAME -r https://github.com/deepseek-ai/DeepGEMM.git -b v0.1.0
EOF
}

# Parse CLI options
if ! OPTIONS=$(getopt -o p:r:b:d:fh --long project_path:,repo:,branch:,dir:,force-clone,help -n "$SCRIPT_NAME" -- "$@"); then
	log_error "Failed to parse command-line options"
	show_usage
	exit 2
fi
eval set -- "$OPTIONS"
log_info "Parsing command-line options: $OPTIONS"

while true; do
	case "$1" in
	-p | --project_path)
		PROJECT_PATH="$2"
		shift 2
		;;
	-r | --repo)
		DEEPGEMM_REPO="$2"
		shift 2
		;;
	-b | --branch)
		DEEPGEMM_BRANCH="$2"
		shift 2
		;;
	-d | --dir)
		INSTALL_DIR="$2"
		shift 2
		;;
	-f | --force-clone)
		FORCE_CLONE=true
		shift
		;;
	-h | --help)
		show_usage
		exit 0
		;;
	--)
		shift
		break
		;;
	*) break ;;
	esac
done

if [[ $# -gt 0 ]]; then
	log_info "Additional arguments provided (ignored): $*"
fi

# Validate project path
if [[ ! -d $PROJECT_PATH ]]; then
	log_error "Project path does not exist: $PROJECT_PATH"
	exit 1
fi

log_info "Project path: $PROJECT_PATH"

# Detect architecture and GPUs (informational; DeepGEMM is CUDA-centric)
detect_arch() { uname -m; }
detect_nvidia_gpu() { command -v nvidia-smi >/dev/null 2>&1 && nvidia-smi -L 2>/dev/null | wc -l || echo 0; }
detect_amd_gpu() { command -v rocm-smi >/dev/null 2>&1 && rocm-smi --showproductname 2>/dev/null | grep -ci 'amd' || echo 0; }

ARCH="$(detect_arch)"
NV_COUNT="$(detect_nvidia_gpu)"
AMD_COUNT="$(detect_amd_gpu)"
log_info "Detected architecture: $ARCH; NVIDIA GPUs: $NV_COUNT; AMD GPUs: $AMD_COUNT"

if [[ $AMD_COUNT -gt 0 && $NV_COUNT -eq 0 ]]; then
	log_warn "AMD GPUs detected without NVIDIA. DeepGEMM is likely CUDA-only; installation may fail."
fi

# Resolve source directory for building/installing
PREFERRED_LOCAL_DIR="$PROJECT_PATH/apps/DeepGEMM"
DEFAULT_CLONE_DIR="$PROJECT_PATH/DeepGEMM"

resolve_install_dir() {
	if [[ -n $INSTALL_DIR ]]; then
		echo "$INSTALL_DIR"
		return 0
	fi
	if [[ -d $PREFERRED_LOCAL_DIR && $FORCE_CLONE == false ]]; then
		echo "$PREFERRED_LOCAL_DIR"
	else
		echo "$DEFAULT_CLONE_DIR"
	fi
}

INSTALL_ROOT="$(resolve_install_dir)"
log_info "Install source directory: $INSTALL_ROOT"

ensure_repo_present() {
	local dir="$1"
	if [[ -d "$dir/.git" ]] || [[ -f "$dir/install.sh" ]] || [[ -f "$dir/pyproject.toml" ]]; then
		log_info "Using existing DeepGEMM source at: $dir"
		return 0
	fi

	log_info "Cloning DeepGEMM from $DEEPGEMM_REPO (branch: $DEEPGEMM_BRANCH) into $dir"
	rm -rf "$dir"
	if ! git clone --recursive --branch "$DEEPGEMM_BRANCH" "$DEEPGEMM_REPO" "$dir"; then
		log_error "Failed to clone DeepGEMM"
		return 1
	fi
	return 0
}

run_install_script() {
	local dir="$1"
	if [[ ! -x "$dir/install.sh" ]]; then
		# Try to make it executable if present
		if [[ -f "$dir/install.sh" ]]; then
			chmod +x "$dir/install.sh" || true
		fi
	fi

	if [[ -x "$dir/install.sh" ]]; then
		log_info "Running upstream install.sh"
		(source .venv/bin/activate && cd "$dir" && ./install.sh)
		return 0
	fi

	# Fallback: try editable pip install; prefer uv if available
	log_warn "install.sh not found or not executable. Falling back to pip install -e ."
	if command -v uv >/dev/null 2>&1; then
		(cd "$dir" && uv run pip install -e .)
	else
		(cd "$dir" && python3 -m pip install -e .)
	fi
}

verify_install() {
	log_info "Verifying DeepGEMM install via Python import"
	if command -v uv >/dev/null 2>&1; then
		uv run python3 - <<'PY'
try:
	import importlib
	import importlib.util
	import pkgutil
	if importlib.util.find_spec("deep_gemm") is None and pkgutil.find_loader("deep_gemm") is None:
		raise ImportError("deep_gemm not importable")
	import deep_gemm  # noqa: F401
	print("OK: deep_gemm imported")
except Exception as e:
	import sys
	print(f"ERROR: {e}", file=sys.stderr)
	raise
PY
	else
		python3 - <<'PY'
try:
	import importlib
	import importlib.util
	import pkgutil
	if importlib.util.find_spec("deep_gemm") is None and pkgutil.find_loader("deep_gemm") is None:
		raise ImportError("deep_gemm not importable")
	import deep_gemm  # noqa: F401
	print("OK: deep_gemm imported")
except Exception as e:
	import sys
	print(f"ERROR: {e}", file=sys.stderr)
	raise
PY
	fi
}

reinstall_precommit_hooks() {
	if command -v uv >/dev/null 2>&1; then
		if ! uv run pre-commit install; then
			log_warn "pre-commit install failed (non-critical)"
			return 0
		fi
	else
		if command -v pre-commit >/dev/null 2>&1; then
			if ! pre-commit install; then
				log_warn "pre-commit install failed (non-critical)"
			fi
		else
			log_warn "pre-commit not found (non-critical)"
		fi
	fi
}

# --- Main flow ---
log_info "Starting DeepGEMM installation"

ensure_repo_present "$INSTALL_ROOT" || {
	log_error "Source acquisition failed"
	exit 1
}

run_install_script "$INSTALL_ROOT" || {
	log_error "DeepGEMM install script failed"
	exit 1
}

verify_install || {
	log_error "DeepGEMM verification failed"
	exit 1
}

reinstall_precommit_hooks || true

log_info "DeepGEMM installation finished successfully"
exit 0
