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

"""Safe testcase preparation for controlled diagnostic experiments."""

from __future__ import annotations

from collections.abc import Callable, Mapping, Sequence
import json
import os
from pathlib import Path, PurePosixPath
import shutil
import sys
from typing import Any

from ..config import AppConfig, project_root
from ..filesystem import sha256_file
from ..ini import IniError, inspect_ini_value
from ..prefix import (
	PrefixSafetyError,
	cleanup_dedicated_wineserver,
	validate_dedicated_prefix,
)
from ..registry import DesktopRegistrySnapshot
from ..result import ResultDirectory
from ..runner import CommandLogger
from ..testcase import TestIniOverride, TestModification, Testcase
from .cleanup import (
	registry_restore_commands as _registry_restore_commands_impl,
	registry_setup_commands as _registry_setup_commands_impl,
	require_command as _require_command,
	validate_disposable_path as _validate_disposable_path,
)
from .models import ExperimentError


_WINDOWS_GAME_DIRECTORY = r"C:\Program Files (x86)\Pristontale EU"
_DISPOSABLE_DIRECTORY_NAME = "ptdiag-disposable"
_PERSISTENT_STATE_CACHE_SENTINEL = "@ptdiag-persistent"
_PERSISTENT_STATE_CACHE_DIRECTORY = "ptdiag-cache/dxvk-state"


def _resource_copy_paths(
	live_game: Path,
	modification: TestModification,
) -> tuple[Path, Path]:
	if modification.kind not in {"copy-file", "normalize-assbin-textures"}:
		raise ExperimentError(f"Unsupported modification kind: {modification.kind}")
	root = live_game.resolve(strict=True)
	source_relative = _safe_resource_relative_path(modification.source)
	destination_relative = _safe_resource_relative_path(modification.destination)
	source = root.joinpath(*source_relative.parts).resolve(strict=True)
	try:
		source.relative_to(root)
	except ValueError as exc:
		raise ExperimentError("Modification source escapes the live game tree") from exc
	if not source.is_file():
		raise ExperimentError(f"Modification source is not a file: {source}")
	destination = root.joinpath(*destination_relative.parts)
	destination_parent = destination.parent.resolve(strict=True)
	try:
		destination_parent.relative_to(root)
	except ValueError as exc:
		raise ExperimentError("Modification destination escapes the live game tree") from exc
	if modification.kind == "copy-file" and (destination.exists() or destination.is_symlink()):
		raise ExperimentError(
			f"Live modification destination already exists; refusing experiment: {destination}"
		)
	if modification.kind == "normalize-assbin-textures":
		if source_relative != destination_relative:
			raise ExperimentError(
				"Assbin normalization source and destination must be the same disposable file"
			)
		if not modification.texture_indices:
			raise ExperimentError("Assbin normalization requires texture_indices")
		if not destination.is_file():
			raise ExperimentError(f"Assbin normalization target is not a file: {destination}")
	return source, destination


def _safe_resource_relative_path(value: str) -> PurePosixPath:
	if "\\" in value:
		raise ExperimentError("Modification paths must use forward slashes")
	path = PurePosixPath(value)
	if path.is_absolute() or not path.parts or any(part in {"", ".", ".."} for part in path.parts):
		raise ExperimentError(f"Unsafe relative resource path: {value!r}")
	return path


def _prepare_disposable_game(
	config: AppConfig,
	result: ResultDirectory,
	logger: CommandLogger,
	modification: TestModification,
	disposable_game: Path,
	source: Path,
) -> dict[str, Any]:
	staging_root = (config.paths.app_root / _DISPOSABLE_DIRECTORY_NAME).absolute()
	_validate_disposable_path(staging_root, disposable_game, result.path.name)
	if disposable_game.exists() or disposable_game.is_symlink():
		raise ExperimentError(f"Disposable game path already exists: {disposable_game}")
	usage = shutil.disk_usage(staging_root.parent)
	copy_bytes = _allocated_tree_bytes(config.paths.live_game)
	if usage.free < copy_bytes + 2 * 1024**3:
		raise ExperimentError(
			f"Insufficient free space for disposable game copy: "
			f"need at least {copy_bytes + 2 * 1024**3}, have {usage.free}"
		)
	_require_command(
		logger.run(
			["mkdir", "-p", "--", str(staging_root)],
			timeout_seconds=15.0,
		),
		"create disposable staging directory",
	)
	_require_command(
		logger.run(
			[
				"cp",
				"-a",
				"--reflink=auto",
				"--sparse=always",
				f"{config.paths.live_game}/.",
				str(disposable_game),
			],
			timeout_seconds=900.0,
		),
		"copy disposable game tree",
	)
	destination_relative = _safe_resource_relative_path(modification.destination)
	destination = disposable_game.joinpath(*destination_relative.parts)
	if modification.kind == "copy-file":
		if destination.exists() or destination.is_symlink():
			raise ExperimentError(
				f"Disposable destination unexpectedly existed before copy: {destination}"
			)
		_require_command(
			logger.run(
				["cp", "--preserve=mode,timestamps", "--", str(source), str(destination)],
				timeout_seconds=30.0,
			),
			"copy selected resource",
		)
		normalization: dict[str, Any] | None = None
		destination_existed_before = False
	elif modification.kind == "normalize-assbin-textures":
		if not destination.is_file() or destination.is_symlink():
			raise ExperimentError(
				f"Disposable Assbin target is missing or unsafe: {destination}"
			)
		if sha256_file(destination) != sha256_file(source):
			raise ExperimentError("Disposable Assbin differs from live source before transform")
		temporary_output = destination.with_name(
			f".{destination.name}.{result.path.name}.normalized"
		)
		if temporary_output.exists() or temporary_output.is_symlink():
			raise ExperimentError(f"Normalization temporary output exists: {temporary_output}")
		normalize_result = logger.run(
			[
				sys.executable,
				"-m",
				"ptdiag",
				"scene",
				"normalize-textures",
				str(destination),
				str(temporary_output),
				*(str(index) for index in modification.texture_indices),
			],
			cwd=project_root(),
			timeout_seconds=240.0,
		)
		_require_command(normalize_result, "normalize embedded Assbin textures")
		try:
			normalization = json.loads(normalize_result.stdout)
		except json.JSONDecodeError as exc:
			raise ExperimentError("Normalizer returned invalid JSON") from exc
		if not isinstance(normalization, dict) or not normalization.get("ok"):
			raise ExperimentError("Normalizer did not report a verified result")
		_require_command(
			logger.run(
				["mv", "--", str(temporary_output), str(destination)],
				timeout_seconds=30.0,
			),
			"install normalized Assbin in disposable tree",
		)
		destination_existed_before = True
	else:
		raise ExperimentError(f"Unsupported modification kind: {modification.kind}")
	destination_hash = sha256_file(destination)
	source_hash = sha256_file(source)
	if modification.kind == "copy-file" and destination_hash != source_hash:
		raise ExperimentError("Copied resource hash differs from its source")
	if modification.kind == "normalize-assbin-textures" and destination_hash == source_hash:
		raise ExperimentError("Normalized Assbin hash did not change")
	ini_changes = _apply_disposable_ini_overrides(
		disposable_game,
		modification.ini_overrides,
		logger,
	)
	return {
		"staging_root": str(staging_root),
		"free_bytes_before": usage.free,
		"source_allocated_bytes": copy_bytes,
		"destination_path": str(destination),
		"destination_existed_before": destination_existed_before,
		"destination_sha256_after": destination_hash,
		"destination_size_after": destination.stat().st_size,
		"normalization": normalization,
		"ini_changes": ini_changes,
	}


def _inspect_ini_overrides(
	game_root: Path,
	overrides: Sequence[TestIniOverride],
) -> list[dict[str, Any]]:
	root = game_root.resolve(strict=True)
	result: list[dict[str, Any]] = []
	for override in overrides:
		path = _safe_game_file(root, override.path)
		state = inspect_ini_value(path, override.section, override.key)
		result.append(
			{
				"path": str(path),
				"section": override.section,
				"key": override.key,
				"expected": override.expected,
				"value": override.value,
				"current": state["value"],
				"matches_expected": state["value"] == override.expected,
				"matches_value": state["value"] == override.value,
				"sha256": state["sha256"],
				"size": state["size"],
				"mtime_ns": state["mtime_ns"],
			}
		)
	return result


def _apply_disposable_ini_overrides(
	disposable_game: Path,
	overrides: Sequence[TestIniOverride],
	logger: CommandLogger,
) -> list[dict[str, Any]]:
	root = disposable_game.resolve(strict=True)
	result: list[dict[str, Any]] = []
	for override in overrides:
		path = _safe_game_file(root, override.path)
		outcome = logger.run(
			_ini_override_command(path, override),
			cwd=project_root(),
			timeout_seconds=30.0,
		)
		_require_command(outcome, f"override [{override.section}] {override.key}")
		try:
			payload = json.loads(outcome.stdout)
		except json.JSONDecodeError as exc:
			raise ExperimentError("INI override returned invalid JSON") from exc
		if not isinstance(payload, dict) or not payload.get("ok"):
			raise ExperimentError("INI override did not report a verified result")
		result.append(payload)
	return result


def _safe_game_file(root: Path, relative: str) -> Path:
	parts = _safe_resource_relative_path(relative).parts
	path = root.joinpath(*parts)
	resolved = path.resolve(strict=True)
	try:
		resolved.relative_to(root)
	except ValueError as exc:
		raise ExperimentError(f"INI path escapes game tree: {relative}") from exc
	if path.is_symlink() or not resolved.is_file():
		raise ExperimentError(f"INI path is missing or unsafe: {path}")
	return resolved


def _ini_override_command(path: Path, override: TestIniOverride) -> list[str]:
	return [
		sys.executable,
		"-m",
		"ptdiag",
		"ini",
		"override",
		str(path),
		override.section,
		override.key,
		override.expected,
		override.value,
	]


def _allocated_tree_bytes(root: Path) -> int:
	total = 0
	for directory, _, filenames in os.walk(root, followlinks=False):
		for filename in filenames:
			path = Path(directory) / filename
			try:
				total += path.lstat().st_blocks * 512
			except FileNotFoundError:
				continue
	return total


def _log_modification_plan(
	logger: CommandLogger,
	config: AppConfig,
	disposable_game: Path,
	source: Path,
	modification: TestModification,
) -> None:
	staging_root = disposable_game.parent
	destination = disposable_game.joinpath(
		*_safe_resource_relative_path(modification.destination).parts
	)
	prefix_link = config.proton7.prefix / "drive_c/Program Files (x86)/Pristontale EU"
	commands: list[tuple[list[str], float]] = [
		(["mkdir", "-p", "--", str(staging_root)], 15.0),
		(
			[
				"cp",
				"-a",
				"--reflink=auto",
				"--sparse=always",
				f"{config.paths.live_game}/.",
				str(disposable_game),
			],
			900.0,
		),
	]
	if modification.kind == "copy-file":
		commands.append(
			(["cp", "--preserve=mode,timestamps", "--", str(source), str(destination)], 30.0)
		)
	elif modification.kind == "normalize-assbin-textures":
		temporary_output = destination.with_name(f".{destination.name}.planned.normalized")
		commands.extend(
			(
				(
					[
						sys.executable,
						"-m",
						"ptdiag",
						"scene",
						"normalize-textures",
						str(destination),
						str(temporary_output),
						*(str(index) for index in modification.texture_indices),
					],
					240.0,
				),
				(["mv", "--", str(temporary_output), str(destination)], 30.0),
			)
		)
	else:
		raise ExperimentError(f"Unsupported modification kind: {modification.kind}")
	for override in modification.ini_overrides:
		ini_path = disposable_game.joinpath(
			*_safe_resource_relative_path(override.path).parts
		)
		commands.append((_ini_override_command(ini_path, override), 30.0))
	commands.extend(
		(
			(["ln", "-sfn", "--", str(disposable_game), str(prefix_link)], 15.0),
			(["ln", "-sfn", "--", os.readlink(prefix_link), str(prefix_link)], 15.0),
			(["rm", "-rf", "--one-file-system", "--", str(disposable_game)], 900.0),
		)
	)
	for command, timeout in commands:
		logger.run(command, timeout_seconds=timeout, dry_run=True)


def _preflight(
	config: AppConfig,
	testcase: Testcase,
	environment: Mapping[str, Any],
	hashes: list[dict[str, Any]],
	*,
	game_processes: Callable[[], list[int]],
) -> list[str]:
	errors: list[str] = []
	if not environment.get("ok"):
		errors.append("Required environment checks failed")
	if any(not item.get("ok", False) for item in hashes if item.get("required", False)):
		errors.append("A required baseline hash is missing or mismatched")
	if testcase.runner != "proton7":
		errors.append(f"Unsupported runner: {testcase.runner}")
	if testcase.game.args != config.game.args:
		errors.append("Testcase game arguments differ from the known PTEU_REBORN baseline")
	if testcase.game.cwd != "live_game":
		errors.append("Only the verified live_game working directory is supported")
	if testcase.display.size != config.game.virtual_desktop:
		errors.append("Testcase desktop size differs from the 1600x900 baseline")
	if testcase.capture.strace_files and shutil.which("strace") is None:
		errors.append("scene-strace requires strace")
	if not testcase.watch.require_new_native_log:
		errors.append("Runtime tests must require a current-run native log")
	try:
		prefix = validate_dedicated_prefix(config)
	except (OSError, ValueError, PrefixSafetyError) as exc:
		errors.append(str(exc))
	else:
		linked_game = prefix / "drive_c/Program Files (x86)/Pristontale EU"
		try:
			if linked_game.resolve(strict=True) != config.paths.live_game.resolve(strict=True):
				errors.append("Dedicated prefix game link does not resolve to live_game")
		except OSError as exc:
			errors.append(f"Dedicated prefix game link is invalid: {exc}")
	active = game_processes()
	if active:
		errors.append(f"Game.exe is already running: {active}")
	return errors


def _dry_run_commands(
	config: AppConfig,
	testcase: Testcase,
	logger: CommandLogger,
	registry_before: DesktopRegistrySnapshot,
	launch_argv: list[str],
	env_overrides: Mapping[str, str],
) -> None:
	cleanup_dedicated_wineserver(config, logger, timeout_seconds=15.0, dry_run=True)
	if testcase.display.virtual_desktop:
		for command in _registry_setup_commands(
			config, registry_before.desktop_name, testcase.display.size
		):
			logger.run(
				command,
				cwd=config.paths.live_game,
				env_overrides=_proton_environment(config),
				timeout_seconds=15.0,
				dry_run=True,
			)
	logger.run(
		launch_argv,
		cwd=config.paths.live_game,
		env_overrides=env_overrides,
		timeout_seconds=testcase.timeout_seconds,
		dry_run=True,
	)
	cleanup_dedicated_wineserver(config, logger, timeout_seconds=15.0, dry_run=True)
	if testcase.display.virtual_desktop:
		for command in _registry_restore_commands(config, registry_before):
			logger.run(
				command,
				cwd=config.paths.live_game,
				env_overrides=_proton_environment(config),
				timeout_seconds=15.0,
				dry_run=True,
			)


def _launch_argv(config: AppConfig, testcase: Testcase, result_directory: Path) -> list[str]:
	game = f"{_WINDOWS_GAME_DIRECTORY}\\{config.game.exe}"
	proton = _proton_command(config, "runinprefix", game, *testcase.game.args)
	if not testcase.capture.strace_files:
		return proton
	raw_directory = result_directory / "raw-strace"
	raw_directory.mkdir(exist_ok=False)
	strace = shutil.which("strace") or "/usr/bin/strace"
	return [
		strace,
		"-ff",
		"-qq",
		"-ttt",
		"-s",
		"1024",
		"-e",
		"trace=%file",
		"-o",
		str(raw_directory / "trace"),
		*proton,
	]


def _proton_command(config: AppConfig, *arguments: str) -> list[str]:
	shim = project_root() / "ptlinux/proton_compat.py"
	return [
		sys.executable,
		str(shim),
		str(config.proton7.proton),
		*arguments,
	]


def _runtime_environment(
	config: AppConfig,
	testcase: Testcase,
	result_directory: Path,
) -> dict[str, str]:
	environment = _proton_environment(config)
	environment.update(testcase.env)
	if environment.get("DXVK_STATE_CACHE_PATH") == _PERSISTENT_STATE_CACHE_SENTINEL:
		environment["DXVK_STATE_CACHE_PATH"] = str(
			(config.paths.app_root / _PERSISTENT_STATE_CACHE_DIRECTORY).absolute()
		)
	environment.update(
		{
			"DXVK_LOG_PATH": str(result_directory),
			"PROTON_LOG": "0",
			"STEAM_COMPAT_INSTALL_PATH": str(config.paths.live_game),
			"SteamAppId": "0",
			"SteamGameId": "0",
		}
	)
	return environment


def _state_cache_plan(
	config: AppConfig,
	environment: Mapping[str, str],
) -> dict[str, Any] | None:
	configured = environment.get("DXVK_STATE_CACHE_PATH")
	if configured is None:
		return None
	cache_directory = Path(configured).expanduser().absolute()
	cache_file = cache_directory / f"{Path(config.game.exe).stem}.dxvk-cache"
	return {
		"directory": str(cache_directory),
		"cache_file": str(cache_file),
		"persistent": True,
		"seed_source": str(config.paths.live_game / cache_file.name),
		"seeded": False,
		"before": _state_cache_file_snapshot(cache_file),
		"after": None,
	}


def _prepare_state_cache(
	config: AppConfig,
	environment: Mapping[str, str],
	logger: CommandLogger,
) -> dict[str, Any]:
	plan = _state_cache_plan(config, environment)
	if plan is None:
		raise ExperimentError("Persistent DXVK state-cache plan is missing")
	cache_directory = Path(str(plan["directory"]))
	expected = (config.paths.app_root / _PERSISTENT_STATE_CACHE_DIRECTORY).absolute()
	if cache_directory != expected:
		raise ExperimentError(
			f"Refusing unmanaged DXVK state-cache directory: {cache_directory}; "
			f"expected {expected}"
		)
	_require_command(
		logger.run(
			["mkdir", "-p", "--", str(cache_directory)],
			timeout_seconds=15.0,
		),
		"create managed DXVK state-cache directory",
	)
	cache_file = Path(str(plan["cache_file"]))
	seed_source = Path(str(plan["seed_source"]))
	if not cache_file.exists() and seed_source.is_file():
		_require_command(
			logger.run(
				[
					"cp",
					"--preserve=mode,timestamps",
					"--",
					str(seed_source),
					str(cache_file),
				],
				timeout_seconds=30.0,
			),
			"seed managed DXVK state cache",
		)
		plan["seeded"] = True
	plan["before"] = _state_cache_file_snapshot(cache_file)
	return plan


def _state_cache_file_snapshot(path: Path) -> dict[str, Any]:
	result: dict[str, Any] = {
		"path": str(path),
		"exists": path.is_file(),
		"size": None,
		"mtime_ns": None,
		"sha256": None,
	}
	if not result["exists"]:
		return result
	try:
		state = path.stat()
		result.update(
			{
				"size": state.st_size,
				"mtime_ns": state.st_mtime_ns,
				"sha256": sha256_file(path),
			}
		)
	except OSError as exc:
		result["error"] = str(exc)
	return result


def _proton_environment(config: AppConfig) -> dict[str, str]:
	return {
		"STEAM_COMPAT_CLIENT_INSTALL_PATH": str(config.proton7.steam_root),
		"STEAM_COMPAT_DATA_PATH": str(config.proton7.compat),
		"WINEPREFIX": str(config.proton7.prefix),
		"WINEDEBUG": "-all",
	}


def _registry_setup_commands(
	config: AppConfig,
	desktop_name: str,
	desktop_size: str,
) -> list[list[str]]:
	return _registry_setup_commands_impl(
		config,
		desktop_name,
		desktop_size,
		proton_command=_proton_command,
	)


def _registry_restore_commands(
	config: AppConfig,
	snapshot: DesktopRegistrySnapshot,
) -> list[list[str]]:
	return _registry_restore_commands_impl(
		config,
		snapshot,
		proton_command=_proton_command,
	)