transcription/tools/run_destructive_tests.py

from __future__ import annotations

import argparse
import ctypes
import os
from ctypes import wintypes
from datetime import datetime
from pathlib import Path
import shlex
import shutil
import subprocess
import sys


RETRY_CANCEL_CHOICES = {"n", "no", "c", "cancel", "a", "abort", "q", "quit"}


def show_phase(title: str) -> None:
    print()
    print(f"========== {title} ==========")


def test_file_unlocked(path: Path) -> bool:
    if not path.exists():
        return True

    if os.name == "nt":
        return _test_file_unlocked_windows(path)

    return _test_file_unlocked_posix(path)


def _test_file_unlocked_windows(path: Path) -> bool:
    generic_read = 0x80000000
    generic_write = 0x40000000
    open_existing = 3
    file_attribute_normal = 0x80
    invalid_handle_value = wintypes.HANDLE(-1).value

    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
    kernel32.CreateFileW.argtypes = [
        wintypes.LPCWSTR,
        wintypes.DWORD,
        wintypes.DWORD,
        wintypes.LPVOID,
        wintypes.DWORD,
        wintypes.DWORD,
        wintypes.HANDLE,
    ]
    kernel32.CreateFileW.restype = wintypes.HANDLE
    kernel32.CloseHandle.argtypes = [wintypes.HANDLE]
    kernel32.CloseHandle.restype = wintypes.BOOL

    handle = kernel32.CreateFileW(
        str(path),
        generic_read | generic_write,
        0,
        None,
        open_existing,
        file_attribute_normal,
        None,
    )
    if handle == invalid_handle_value:
        return False

    kernel32.CloseHandle(handle)
    return True


def _test_file_unlocked_posix(path: Path) -> bool:
    import fcntl

    fd = os.open(path, os.O_RDWR)
    try:
        try:
            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
        except BlockingIOError:
            return False
        else:
            fcntl.flock(fd, fcntl.LOCK_UN)
            return True
    finally:
        os.close(fd)


def wait_for_restore_preflight(db_file_path: Path) -> bool:
    attempt = 1
    while not test_file_unlocked(db_file_path):
        show_phase("Restore Preflight")
        print(f"WARNING: Restore preflight blocked: database appears to be in use: {db_file_path}")
        print("Close conflicting applications (for example DB Browser for SQLite, uvicorn, or any process using this DB).")
        print("Restore is paused and waiting for your input.")
        answer = input(
            f"Attempt {attempt}: type 'retry' to check again, or 'cancel' to skip restore: "
        ).strip()
        print(f"Input received: '{answer}'")
        if answer.lower() in RETRY_CANCEL_CHOICES:
            return False
        attempt += 1
        print("Re-checking database lock now...")

    if attempt > 1:
        print("Restore preflight passed: database lock released.")
    return True


def restore_backup(backup_path: Path, data_path: Path) -> None:
    if not backup_path.exists():
        raise FileNotFoundError(f"Backup path not found: {backup_path}")

    if data_path.exists():
        shutil.rmtree(data_path)

    shutil.copytree(backup_path, data_path)
    print(f"Restored data from backup: {backup_path}")


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Run potentially destructive tests with backup/restore protection."
    )
    parser.add_argument("--auto-restore", action="store_true", help="Restore immediately after successful tests.")
    parser.add_argument("--keep-backup", action="store_true", help="Keep the backup even after a successful restore.")
    parser.add_argument(
        "--skip-restore-prompt",
        action="store_true",
        help="Do not prompt after successful tests; keep the backup for a later manual restore.",
    )
    parser.add_argument(
        "--restore-from",
        help="Restore from an existing backup name or absolute backup path instead of running tests.",
    )
    parser.add_argument("command", nargs=argparse.REMAINDER, help="Command to run after '--'.")
    args = parser.parse_args(argv)

    if args.restore_from and args.command:
        parser.error("--restore-from cannot be combined with a test command.")

    if not args.restore_from and not args.command:
        parser.error("A test command is required unless --restore-from is provided.")

    return args


def normalize_command(command: list[str]) -> list[str]:
    if command and command[0] == "--":
        return command[1:]
    return command


def run_command(command: list[str]) -> int:
    show_phase("Test Phase")
    print(f"Running command: {shlex.join(command)}")
    try:
        completed = subprocess.run(command, check=False)
    except OSError as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    return completed.returncode


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv or sys.argv[1:])
    repo_root = Path(__file__).resolve().parent.parent
    data_dir = repo_root / "data"
    backup_root = repo_root / ".test-backups"
    db_path = data_dir / "transcription.db"

    if not data_dir.exists():
        raise FileNotFoundError(f"Data directory not found: {data_dir}")

    backup_root.mkdir(parents=True, exist_ok=True)

    if args.restore_from:
        restore_path = Path(args.restore_from)
        if not restore_path.is_absolute():
            restore_path = backup_root / restore_path

        show_phase("Restore Phase")
        if not wait_for_restore_preflight(db_path):
            print(f"Restore cancelled. Backup preserved at: {restore_path}")
            return 1

        restore_backup(restore_path, data_dir)
        return 0

    command = normalize_command(list(args.command))
    if not command:
        raise ValueError("No test command provided.")

    show_phase("Backup Phase")
    if not test_file_unlocked(db_path):
        print(f"WARNING: Backup preflight warning: database appears to be in use: {db_path}")
        print("WARNING: Proceeding with backup, but hot backups can capture an in-flight state.")

    timestamp = datetime.now().strftime("%Y%m%d-%H%M%S")
    backup_name = f"data-backup-{timestamp}"
    backup_path = backup_root / backup_name
    shutil.copytree(data_dir, backup_path)
    print(f"Created backup: {backup_path}")

    test_exit_code = run_command(command)
    if test_exit_code != 0:
        show_phase("Post-Test")
        print(f"Tests failed (exit code {test_exit_code}).")
        print(f"Backup preserved at: {backup_path}")
        print("Inspect state, then restore manually if needed:")
        print(f"  uv run python tools/run_destructive_tests.py --restore-from {backup_name}")
        return test_exit_code

    should_restore = False
    if args.auto_restore:
        should_restore = True
    elif args.skip_restore_prompt:
        should_restore = False
    else:
        answer = input("Tests passed. Restore data backup now? [y/N] ").strip().lower()
        if answer in {"y", "yes"}:
            should_restore = True

    if not should_restore:
        print(f"Restore skipped by user. Backup kept at: {backup_path}")
        print("Restore later with:")
        print(f"  uv run python tools/run_destructive_tests.py --restore-from {backup_name}")
        return 0

    show_phase("Restore Phase")
    if not wait_for_restore_preflight(db_path):
        print(f"Restore cancelled. Backup preserved at: {backup_path}")
        return 1

    try:
        restore_backup(backup_path, data_dir)
        if args.keep_backup:
            print(f"Kept backup: {backup_path}")
        else:
            shutil.rmtree(backup_path)
            print(f"Deleted backup: {backup_path}")
    except Exception as exc:
        print("WARNING: Restore failed. Your current data remains unchanged.")
        print(f"WARNING: {exc}")
        print("Likely cause: another process has data/transcription.db open.")
        print("Stop the process and retry restore with:")
        print(f"  uv run python tools/run_destructive_tests.py --restore-from {backup_name}")
        return 1

    return 0


if __name__ == "__main__":
    raise SystemExit(main())