#!/usr/bin/env python3
"""Offline checkpoint completion/lock contracts; no live servers or transfers.

The loader tests execute its production classification, MD5 preflight and replay
blocks in a temporary environment. They do not claim whole-script/service tests.
"""
import json
import os
from pathlib import Path
import shutil
import subprocess
import sys
import tempfile
import unittest


BIN = Path(__file__).resolve().parents[1] / "Server/Unix/p4/common/bin"
BASH = os.environ.get("SDP_TEST_BASH", shutil.which("bash"))


class RecoveryFixture(unittest.TestCase):
    def setUp(self):
        temporary = tempfile.TemporaryDirectory(prefix="sdp-recovery-completion-")
        self.addCleanup(temporary.cleanup)
        self.root = Path(temporary.name).resolve()
        self.logs = self.root / "logs"
        self.checkpoints = self.root / "checkpoints.gz-parent"
        self.database = self.root / "database"
        self.offline = self.root / "offline"
        for directory in (self.logs, self.checkpoints, self.database, self.offline):
            directory.mkdir()
        self.marker = self.logs / "ckp_running.txt"
        self.calls = self.root / "p4d-calls.jsonl"
        self.p4d = self.root / "p4d-fixture"
        self.p4d.write_text("#!" + sys.executable + r'''
import json, os, pathlib, sys
args = sys.argv[1:]
with open(os.environ["TEST_CALLS"], "a") as output:
    output.write(json.dumps(args) + "\n")
if not ("-jr" in args or "-jrp" in args):
    print("unexpected validation/capability prepass", file=sys.stderr)
    sys.exit(88)
target = pathlib.Path(args[args.index("-r") + 1])
if os.environ.get("TEST_RESTORE_FAILURE"):
    (target / "db.partial").write_text("incomplete restore")
    print("injected actual restore failure", file=sys.stderr)
    sys.exit(7)
(target / "db.counters").write_text("restored exact fixture counters")
''')
        self.p4d.chmod(0o755)
        self.env = dict(os.environ, TEST_ROOT=str(self.root), LOGS=str(self.logs),
                        LOGFILE=str(self.logs / "recovery.log"), TEST_CALLS=str(self.calls),
                        P4CBIN=str(self.root), P4DBIN=str(self.p4d),
                        P4ROOT=str(self.database), CHECKPOINTS=str(self.checkpoints),
                        P4CKP_PREFIX="p4_test", NumCheckPointThreads="6",
                        BACKUP_FUNCTIONS=str(BIN / "backup_functions.sh"))
        for name in ("CkpLockToken", "CkpLockPid", "CkpLockSubshell", "SNS_ALERT_TOPIC_ARN"):
            self.env.pop(name, None)
        self.prefix = '''source "$BACKUP_FUNCTIONS"
OFFLINE_DB="$TEST_ROOT/offline"
log () { printf '%s\\n' "$*" >> "$LOGFILE"; }
mail_log_file () { return 0; }
sdp_fqdn () { printf fixture-host; }
tty () { return 1; }
'''

    def run_shell(self, script, **extra):
        return subprocess.run([BASH, "-c", self.prefix + script], cwd=self.root,
                              env=dict(self.env, **extra), text=True,
                              capture_output=True, timeout=10)

    def recorded_calls(self):
        return [json.loads(line) for line in self.calls.read_text().splitlines()] if self.calls.exists() else []

    def checkpoint(self, kind="gzip", sidecar="normal"):
        source = self.checkpoints / "p4_test.ckp.1"
        if kind == "parallel":
            source.mkdir()
            (source / "db.config.ckp.gz").write_text("fixture checkpoint segment")
        else:
            if kind == "gzip":
                source = source.with_suffix(source.suffix + ".gz")
            source.write_text("fixture checkpoint content")
        if sidecar != "missing":
            path = Path(str(source) + ".md5")
            if sidecar == "stem":
                path = Path(str(source)[:-3] + ".md5")
            if sidecar == "directory":
                path.mkdir()
            else:
                path.write_text("Existing completion sidecar; no separate hash pass.\n")
        return source


class FailureLockRelease(RecoveryFixture):
    def test_die_releases_owned_marker_and_keeps_original_failure_status(self):
        result = self.run_shell('ckp_running; die "fixture operation failed"')
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertFalse(self.marker.exists())
        self.assertIn("fixture operation failed", (self.logs / "recovery.log").read_text())

    def test_die_preserves_unowned_legacy_marker(self):
        self.marker.write_text("other owner\n")
        result = self.run_shell('die "failure without ownership"')
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual(self.marker.read_text(), "other owner\n")

    def test_die_preserves_replaced_marker(self):
        result = self.run_shell('''ckp_running
printf 'replacement owner\\n' > "$LOGS/ckp_running.txt"
die "operation failed after ownership changed"
''')
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual(self.marker.read_text(), "replacement owner\n")

    def test_child_die_does_not_release_parent_marker(self):
        result = self.run_shell('''ckp_running
( die "child failure" ); status=$?
[[ "$status" == 1 && -s "$LOGS/ckp_running.txt" ]] || exit 71
ckp_complete
''')
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertFalse(self.marker.exists())

    def test_child_completion_does_not_release_parent_marker(self):
        result = self.run_shell('''ckp_running
( ckp_complete ); status=$?
[[ "$status" == 1 && -s "$LOGS/ckp_running.txt" ]] || exit 71
ckp_complete
''')
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertFalse(self.marker.exists())

    def test_failed_marker_removal_does_not_recurse_or_replace_original_error(self):
        result = self.run_shell('''ckp_running
rm () { echo 'injected marker unlink failure' >&2; return 13; }
die "original restore failure"
''')
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertTrue(self.marker.exists())
        log = (self.logs / "recovery.log").read_text()
        self.assertEqual(log.count("ERROR!!!"), 1)
        self.assertIn("original restore failure", log)
        self.assertIn("injected marker unlink failure", log + result.stderr)

    def test_missing_owned_marker_still_reports_original_failure_once(self):
        result = self.run_shell('ckp_running; rm -- "$LOGS/ckp_running.txt"; die "original failure"')
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertFalse(self.marker.exists())
        self.assertEqual((self.logs / "recovery.log").read_text().count("ERROR!!!"), 1)


class OfflineRestoreCompletion(RecoveryFixture):
    def restore(self, kind, **extra):
        return self.run_shell('''DoParallelCheckpoints=$TEST_PARALLEL
NumCheckPointThreads=6
ckp_running
recreate_offline_db_files
ckp_complete
''', TEST_PARALLEL="1" if kind == "parallel" else "0", **extra)

    def assert_single_restore(self, source, parallel=False):
        calls = self.recorded_calls()
        expected = ["-r", str(self.offline)]
        if parallel:
            expected += ["-z", "-jrp", "-N", "6", str(source)]
        else:
            if source.suffix == ".gz":
                expected += ["-z"]
            expected += ["-jr", str(source)]
        self.assertEqual(calls, [expected])

    def test_serial_restore_uses_completion_sidecar_without_verification_prepass(self):
        source = self.checkpoint()
        result = self.restore("gzip")
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assert_single_restore(source)
        self.assertTrue((self.offline / "offline_db_usable.txt").is_file())

    def test_parallel_restore_does_not_require_jvp_capability(self):
        source = self.checkpoint("parallel")
        result = self.restore("parallel")
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assert_single_restore(source, parallel=True)

    def test_plain_restore_keeps_existing_flags(self):
        source = self.checkpoint("plain")
        result = self.restore("plain")
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assert_single_restore(source)

    def test_gzip_stem_sidecar_is_accepted(self):
        source = self.checkpoint(sidecar="stem")
        result = self.restore("gzip")
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assert_single_restore(source)

    def test_missing_sidecar_preserves_existing_offline_database(self):
        self.checkpoint(sidecar="missing")
        (self.offline / "db.counters").write_text("original database")
        (self.offline / "offline_db_usable.txt").write_text("original usable marker")
        result = self.restore("gzip")
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual((self.offline / "db.counters").read_text(), "original database")
        self.assertTrue((self.offline / "offline_db_usable.txt").exists())
        self.assertEqual(self.recorded_calls(), [])
        self.assertFalse(self.marker.exists())

    def test_missing_checkpoint_does_not_accept_orphan_sidecar(self):
        source = self.checkpoint()
        source.unlink()
        (self.offline / "db.counters").write_text("original database")
        result = self.restore("gzip")
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual((self.offline / "db.counters").read_text(), "original database")
        self.assertEqual(self.recorded_calls(), [])

    def test_missing_parallel_sidecar_preserves_existing_offline_database(self):
        self.checkpoint("parallel", sidecar="missing")
        (self.offline / "db.counters").write_text("original database")
        result = self.restore("parallel")
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual((self.offline / "db.counters").read_text(), "original database")
        self.assertEqual(self.recorded_calls(), [])

    def test_directory_named_md5_is_not_selected_as_completion_sidecar(self):
        self.checkpoint(sidecar="directory")
        (self.offline / "db.counters").write_text("original database")
        result = self.restore("gzip")
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual((self.offline / "db.counters").read_text(), "original database")
        self.assertEqual(self.recorded_calls(), [])

    def test_directory_or_fifo_serial_source_does_not_remove_database(self):
        source = self.checkpoint()
        source.unlink()
        (self.offline / "db.counters").write_text("original database")
        (self.offline / "offline_db_usable.txt").write_text("original usability evidence")
        for kind in ("directory", "fifo"):
            with self.subTest(kind=kind):
                source.mkdir() if kind == "directory" else os.mkfifo(source)
                result = self.restore("gzip")
                self.assertEqual(result.returncode, 1, result.stderr)
                self.assertEqual((self.offline / "db.counters").read_text(), "original database")
                self.assertEqual((self.offline / "offline_db_usable.txt").read_text(), "original usability evidence")
                self.assertEqual(self.recorded_calls(), [])
                source.rmdir() if kind == "directory" else source.unlink()

    def test_restore_failure_is_not_reported_as_usable_or_successful(self):
        source = self.checkpoint()
        (self.offline / "offline_db_usable.txt").write_text("old usable marker")
        result = self.restore("gzip", TEST_RESTORE_FAILURE="1")
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assert_single_restore(source)
        self.assertFalse((self.offline / "offline_db_usable.txt").exists())
        self.assertFalse(self.marker.exists())
        self.assertIn("injected actual restore failure", (self.logs / "recovery.log").read_text())


class LoaderCompletionPhases(RecoveryFixture):
    def setUp(self):
        super().setUp()
        source = (BIN / "load_checkpoint.sh").read_text()
        classify_start = source.index('if [[ "$ReplayJournalsOnly" -eq 0 ]]; then', source.index('Phase 1 - Preflight Checks.'))
        classify_end = source.index('# Set or check ServerID value.', classify_start)
        completion_start = source.index('if [[ "$ReplayJournalsOnly" -eq 0 ]]; then', source.index('Unexpected P4ROOT value of:'))
        summary = source.index('Preflight Summary:', completion_start)
        completion_end = source.rfind("\n", completion_start, summary) + 1
        replay_start = source.index('if [[ "$ReplayJournalsOnly" -eq 0 ]]; then', source.index('Phase 4 - Load Checkpoint in P4ROOT.'))
        replay_end = source.index('if [[ "$JournalCount" -gt 0 ]]; then', replay_start)
        self.phases = source[classify_start:classify_end] + source[completion_start:completion_end] + source[replay_start:replay_end]
        selector_start = source.index("function find_latest_checkpoint_md5 () {")
        selector_end = source.index("\n#------------------------------------------------------------------------------", selector_start)
        latest_start = source.index('if [[ "$GetLatestCheckpoint" -eq 1 ]]; then', source.index('# Find the newest checkpoint checksum'))
        latest_end = source.index('# If the user requested to find the latest journals', latest_start)
        self.latest = source[selector_start:selector_end] + source[latest_start:latest_end]

    def load(self, source, **extra):
        return self.run_shell('''msg () { log "$*"; }
bail () { log "$*"; exit 1; }
usage () { bail "$*"; }
ReplayJournalsOnly=0; GetLatestCheckpoint=0; LicenseFileNeeded=0
LoadParallelCheckpoint=0; Threads=4; DoParallelCheckpoints=1; NumCheckPointThreads=6
CompressFlag=; CaseFlag=; ReplicaTypeTag=ha; ExcludedTables=
Checkpoint=$TEST_CHECKPOINT
''' + self.phases, TEST_CHECKPOINT=str(source), **extra)

    def test_serial_and_parallel_load_execute_only_the_normal_restore(self):
        for kind in ("plain", "gzip", "parallel"):
            with self.subTest(kind=kind):
                source = self.checkpoint(kind)
                result = self.load(source)
                self.assertEqual(result.returncode, 0, result.stderr)
                expected = ["-r", str(self.database)]
                expected += ["-N", "6", "-jrp", str(source)] if kind == "parallel" else ["-jr", str(source)]
                self.assertEqual(self.recorded_calls(), [expected])
                self.calls.unlink()
                for path in self.checkpoints.iterdir():
                    shutil.rmtree(path) if path.is_dir() else path.unlink()

    def test_gzip_stem_sidecar_works_when_parent_directory_contains_dot_gz(self):
        source = self.checkpoint(sidecar="stem")
        result = self.load(source)
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertEqual(self.recorded_calls(), [["-r", str(self.database), "-jr", str(source)]])

    def test_missing_sidecar_plain_gzip_and_parallel_are_rejected_before_restore(self):
        for kind in ("plain", "gzip", "parallel"):
            with self.subTest(kind=kind):
                source = self.checkpoint(kind, sidecar="missing")
                result = self.load(source)
                self.assertEqual(result.returncode, 1, result.stderr)
                self.assertEqual(self.recorded_calls(), [])
                self.assertFalse((self.database / "db.counters").exists())
                shutil.rmtree(source) if source.is_dir() else source.unlink()

    def test_directory_named_md5_is_not_a_completion_file(self):
        source = self.checkpoint(sidecar="directory")
        result = self.load(source)
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual(self.recorded_calls(), [])

    def test_unreadable_source_is_rejected_before_restore(self):
        source = self.checkpoint()
        source.unlink()
        source.symlink_to(self.checkpoints / "missing checkpoint")
        result = self.load(source)
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual(self.recorded_calls(), [])

    def test_actual_restore_failure_remains_a_failure(self):
        source = self.checkpoint()
        result = self.load(source, TEST_RESTORE_FAILURE="1")
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual(self.recorded_calls(), [["-r", str(self.database), "-jr", str(source)]])
        self.assertIn("Checkpoint replay to P4ROOT failed", (self.logs / "recovery.log").read_text())

    def select_latest(self):
        return self.run_shell('''msg () { log "$*"; }
bail () { log "$*"; exit 1; }
GetLatestCheckpoint=1; LoadParallelCheckpoint=0
CheckpointsDir=$CHECKPOINTS; JournalPrefix="$CHECKPOINTS/$P4CKP_PREFIX"
''' + self.latest + '\nprintf "%s\\n" "$Checkpoint"\n')

    def test_latest_selector_ignores_directory_and_fifo_checksum_candidates(self):
        source = self.checkpoint()
        os.utime(Path(str(source) + ".md5"), (1, 1))
        (self.checkpoints / "p4_test.ckp.2.md5").mkdir()
        os.mkfifo(self.checkpoints / "p4_test.ckp.3.md5")
        result = self.select_latest()
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertEqual(result.stdout.strip(), str(source))
        self.assertEqual(self.recorded_calls(), [])

    def test_latest_gzip_fallback_preserves_parent_directory_containing_md5(self):
        self.checkpoints = self.root / "checkpoints.md5-parent"
        self.checkpoints.mkdir()
        self.env["CHECKPOINTS"] = str(self.checkpoints)
        source = self.checkpoint(sidecar="stem")
        result = self.select_latest()
        self.assertEqual(result.returncode, 0, result.stderr)
        self.assertEqual(result.stdout.strip(), str(source))

    def test_latest_selector_rejects_fifo_as_serial_checkpoint(self):
        source = self.checkpoint()
        source.unlink()
        os.mkfifo(source)
        result = self.select_latest()
        self.assertEqual(result.returncode, 1, result.stderr)
        self.assertEqual(self.recorded_calls(), [])


if __name__ == "__main__":
    unittest.main()