pristontale-eu-linux/src/ptdiag/experiments/cleanup.py

"""Experiment cleanup, prefix-link rollback, and registry restoration."""

from __future__ import annotations

from collections.abc import Callable, Mapping, Sequence
import os
from pathlib import Path

from ..config import AppConfig
from ..prefix import cleanup_dedicated_wineserver
from ..registry import DesktopRegistrySnapshot
from ..result import ResultDirectory
from ..runner import CommandLogger, CommandResult
from .models import ExperimentError


EXPLORER_KEY = r"HKCU\Software\Wine\Explorer"
DESKTOPS_KEY = r"HKCU\Software\Wine\Explorer\Desktops"
DISPOSABLE_DIRECTORY_NAME = "ptdiag-disposable"

ProtonCommand = Callable[..., list[str]]
WineserverCleaner = Callable[..., CommandResult]


def replace_game_link(
	prefix_link: Path,
	disposable_game: Path,
	expected_live_game: Path,
	logger: CommandLogger,
) -> str:
	"""Point the dedicated prefix at a verified disposable game tree."""

	if not prefix_link.is_symlink():
		raise ExperimentError(
			f"Dedicated-prefix game path is not a symlink: {prefix_link}"
		)
	if prefix_link.resolve(strict=True) != expected_live_game.resolve(strict=True):
		raise ExperimentError("Dedicated-prefix game link changed after source preflight")
	original_target = os.readlink(prefix_link)
	require_command(
		logger.run(
			["ln", "-sfn", "--", str(disposable_game), str(prefix_link)],
			cwd=prefix_link.parent,
			timeout_seconds=15.0,
		),
		"point dedicated prefix at disposable game",
	)
	if (
		not prefix_link.is_symlink()
		or prefix_link.resolve(strict=True) != disposable_game.resolve(strict=True)
	):
		raise ExperimentError(
			"Dedicated-prefix game link did not resolve to disposable copy"
		)
	return original_target


def restore_game_link(
	prefix_link: Path,
	original_target: str,
	expected_live_game: Path,
	logger: CommandLogger,
) -> None:
	"""Restore the dedicated-prefix game link and verify its exact target."""

	require_command(
		logger.run(
			["ln", "-sfn", "--", original_target, str(prefix_link)],
			cwd=prefix_link.parent,
			timeout_seconds=15.0,
		),
		"restore dedicated-prefix game link",
	)
	if not prefix_link.is_symlink():
		raise ExperimentError("Restored dedicated-prefix game path is not a symlink")
	if os.readlink(prefix_link) != original_target:
		raise ExperimentError("Dedicated-prefix game link text was not restored exactly")
	if prefix_link.resolve(strict=True) != expected_live_game.resolve(strict=True):
		raise ExperimentError("Dedicated-prefix game link target was not restored")


def remove_disposable_game(
	config: AppConfig,
	result: ResultDirectory,
	disposable_game: Path,
	logger: CommandLogger,
	*,
	directory_name: str = DISPOSABLE_DIRECTORY_NAME,
) -> None:
	"""Remove only the result-scoped disposable tree beneath its managed root."""

	staging_root = (config.paths.app_root / directory_name).absolute()
	validate_disposable_path(
		staging_root,
		disposable_game,
		result.path.name,
		directory_name=directory_name,
	)
	require_command(
		logger.run(
			["rm", "-rf", "--one-file-system", "--", str(disposable_game)],
			timeout_seconds=900.0,
		),
		"remove disposable game tree",
	)
	if disposable_game.exists() or disposable_game.is_symlink():
		raise ExperimentError(
			f"Disposable game cleanup left its target: {disposable_game}"
		)
	try:
		staging_root.rmdir()
	except OSError:
		pass


def validate_disposable_path(
	staging_root: Path,
	path: Path,
	expected_name: str,
	*,
	directory_name: str = DISPOSABLE_DIRECTORY_NAME,
) -> None:
	"""Reject cleanup targets outside the managed, result-scoped staging root."""

	if path.parent != staging_root or path.name != expected_name:
		raise ExperimentError(f"Unsafe disposable game path: {path}")
	if staging_root.name != directory_name:
		raise ExperimentError(f"Unsafe disposable staging root: {staging_root}")


def require_command(outcome: CommandResult, operation: str) -> None:
	if not command_ok(outcome):
		raise ExperimentError(
			f"Could not {operation} ({outcome.returncode}): "
			f"{outcome.error or outcome.stderr.strip() or 'command failed'}"
		)


def registry_setup_commands(
	config: AppConfig,
	desktop_name: str,
	desktop_size: str,
	*,
	proton_command: ProtonCommand,
) -> list[list[str]]:
	return [
		proton_command(
			config,
			"runinprefix",
			"reg.exe",
			"add",
			EXPLORER_KEY,
			"/v",
			"Desktop",
			"/t",
			"REG_SZ",
			"/d",
			desktop_name,
			"/f",
		),
		proton_command(
			config,
			"runinprefix",
			"reg.exe",
			"add",
			DESKTOPS_KEY,
			"/v",
			desktop_name,
			"/t",
			"REG_SZ",
			"/d",
			desktop_size,
			"/f",
		),
	]


def registry_restore_commands(
	config: AppConfig,
	snapshot: DesktopRegistrySnapshot,
	*,
	proton_command: ProtonCommand,
) -> list[list[str]]:
	return [
		registry_restore_command(
			config,
			EXPLORER_KEY,
			"Desktop",
			snapshot.selected_desktop.exists,
			snapshot.selected_desktop.value,
			proton_command=proton_command,
		),
		registry_restore_command(
			config,
			DESKTOPS_KEY,
			snapshot.desktop_name,
			snapshot.desktop_size.exists,
			snapshot.desktop_size.value,
			proton_command=proton_command,
		),
	]


def registry_restore_command(
	config: AppConfig,
	key: str,
	name: str,
	existed: bool,
	value: str | None,
	*,
	proton_command: ProtonCommand,
) -> list[str]:
	if existed:
		assert value is not None
		return proton_command(
			config,
			"runinprefix",
			"reg.exe",
			"add",
			key,
			"/v",
			name,
			"/t",
			"REG_SZ",
			"/d",
			value,
			"/f",
		)
	return proton_command(
		config,
		"runinprefix",
		"reg.exe",
		"delete",
		key,
		"/v",
		name,
		"/f",
	)


def run_checked_commands(
	logger: CommandLogger,
	commands: Sequence[Sequence[str]],
	*,
	cwd: Path,
	env_overrides: Mapping[str, str],
) -> None:
	for command in commands:
		outcome = logger.run(
			command,
			cwd=cwd,
			env_overrides=env_overrides,
			timeout_seconds=15.0,
		)
		if not command_ok(outcome):
			raise ExperimentError(
				f"Command failed ({outcome.returncode}): {list(command)!r}: "
				f"{outcome.error or outcome.stderr.strip()}"
			)


def cleanup_checked(
	config: AppConfig,
	logger: CommandLogger,
	*,
	cleaner: WineserverCleaner = cleanup_dedicated_wineserver,
) -> None:
	"""Stop the dedicated wineserver while accepting its idle return contract."""

	outcome = cleaner(config, logger, timeout_seconds=15.0)
	if (
		outcome.returncode == 1
		and not outcome.timed_out
		and outcome.error is None
		and not outcome.stdout
		and not outcome.stderr
	):
		return
	if not command_ok(outcome):
		raise ExperimentError(
			"Dedicated-prefix wineserver cleanup failed: "
			f"{outcome.error or outcome.stderr.strip() or outcome.returncode}"
		)


def command_ok(result: CommandResult) -> bool:
	return not result.timed_out and result.error is None and result.returncode == 0