"""Failure and recovery tests for transactional checkpoint writes."""

from __future__ import annotations

import os
from pathlib import Path
from unittest.mock import MagicMock

import pytest
import torch

import obliteratus.abliterate as abliterate
import obliteratus.persistence_contracts as persistence
from obliteratus.abliterate import AbliterationPipeline, _atomic_checkpoint_directory


pytestmark = pytest.mark.cpu


class SimulatedCancellation(BaseException):
    """Cancellation signal used to prove BaseException rollback behavior."""


def _temporary_artifacts(parent: Path, name: str) -> list[Path]:
    return [
        *parent.glob(f".{name}.staging-*"),
        *parent.glob(f".{name}.backup-*"),
    ]


def test_atomic_checkpoint_replaces_existing_destination(tmp_path):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    (destination / "old.txt").write_text("old", encoding="utf-8")

    with _atomic_checkpoint_directory(destination) as staging:
        (staging / "new.txt").write_text("new", encoding="utf-8")

    assert not (destination / "old.txt").exists()
    assert (destination / "new.txt").read_text() == "new"
    assert _temporary_artifacts(tmp_path, destination.name) == []


def test_atomic_checkpoint_replaces_an_invalid_file_destination(tmp_path):
    destination = tmp_path / "checkpoint"
    destination.write_text("not a checkpoint", encoding="utf-8")

    with _atomic_checkpoint_directory(destination) as staging:
        (staging / "config.json").write_text("{}", encoding="utf-8")

    assert destination.is_dir()
    assert (destination / "config.json").read_text() == "{}"
    assert _temporary_artifacts(tmp_path, destination.name) == []


def test_atomic_checkpoint_warns_if_obsolete_backup_cannot_be_removed(
    tmp_path,
    monkeypatch,
    caplog,
):
    destination = tmp_path / "checkpoint"
    destination.write_text("old", encoding="utf-8")
    real_remove = persistence._remove_checkpoint_path

    def fail_backup_cleanup(path):
        if ".backup-" in path.name:
            raise PermissionError("simulated cleanup denial")
        return real_remove(path)

    monkeypatch.setattr(persistence, "_remove_checkpoint_path", fail_backup_cleanup)
    with _atomic_checkpoint_directory(destination) as staging:
        (staging / "config.json").write_text("{}", encoding="utf-8")

    assert (destination / "config.json").is_file()
    assert "could not be removed" in caplog.text
    backups = list(tmp_path.glob(".checkpoint.backup-*"))
    assert len(backups) == 1
    assert backups[0].read_text() == "old"


def test_atomic_checkpoint_preserves_destination_on_write_failure(tmp_path):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    sentinel = destination / "sentinel.txt"
    sentinel.write_text("preserve me", encoding="utf-8")

    with pytest.raises(OSError, match="simulated write failure"):
        with _atomic_checkpoint_directory(destination) as staging:
            (staging / "partial.bin").write_bytes(b"partial")
            raise OSError("simulated write failure")

    assert sentinel.read_text() == "preserve me"
    assert _temporary_artifacts(tmp_path, destination.name) == []


def test_atomic_checkpoint_cleans_staging_on_cancellation_during_write(tmp_path):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    sentinel = destination / "sentinel.txt"
    sentinel.write_text("preserve me", encoding="utf-8")

    with pytest.raises(SimulatedCancellation):
        with _atomic_checkpoint_directory(destination) as staging:
            (staging / "partial.bin").write_bytes(b"partial")
            raise SimulatedCancellation()

    assert sentinel.read_text(encoding="utf-8") == "preserve me"
    assert _temporary_artifacts(tmp_path, destination.name) == []


def test_atomic_checkpoint_preserves_destination_on_file_fsync_failure(
    tmp_path,
    monkeypatch,
):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    sentinel = destination / "sentinel.txt"
    sentinel.write_text("preserve me", encoding="utf-8")

    def fail_sync(_path):
        raise OSError("simulated fsync failure")

    monkeypatch.setattr(persistence, "_sync_file", fail_sync)
    with pytest.raises(OSError, match="simulated fsync failure"):
        with _atomic_checkpoint_directory(destination) as staging:
            (staging / "partial.bin").write_bytes(b"partial")

    assert sentinel.read_text(encoding="utf-8") == "preserve me"
    assert _temporary_artifacts(tmp_path, destination.name) == []


def test_new_checkpoint_parent_sync_failure_rolls_back_then_allows_retry(
    tmp_path,
    monkeypatch,
):
    destination = tmp_path / "checkpoint"
    real_sync_directory = persistence._sync_directory
    sync_calls = 0
    failed = False

    def fail_after_promotion(path):
        nonlocal sync_calls, failed
        sync_calls += 1
        if sync_calls == 2 and not failed:
            failed = True
            raise OSError("simulated parent fsync failure")
        return real_sync_directory(path)

    monkeypatch.setattr(persistence, "_sync_directory", fail_after_promotion)

    with pytest.raises(OSError, match="simulated parent fsync failure"):
        with _atomic_checkpoint_directory(destination) as staging:
            (staging / "new.txt").write_text("first attempt", encoding="utf-8")

    assert not destination.exists()
    assert _temporary_artifacts(tmp_path, destination.name) == []

    with _atomic_checkpoint_directory(destination) as staging:
        (staging / "new.txt").write_text("retry", encoding="utf-8")

    assert (destination / "new.txt").read_text(encoding="utf-8") == "retry"


@pytest.mark.parametrize(
    "failure_phase",
    [
        "tree_sync",
        "backup_replace",
        "backup_parent_sync",
        "promotion_replace",
        "promotion_parent_sync",
    ],
)
def test_atomic_checkpoint_cancellation_restores_and_unlocks_for_retry(
    tmp_path,
    monkeypatch,
    failure_phase,
):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    sentinel = destination / "sentinel.txt"
    sentinel.write_text("preserve me", encoding="utf-8")
    real_replace = os.replace
    real_sync_file = persistence._sync_file
    real_sync_directory = persistence._sync_directory
    replace_calls = 0
    directory_sync_calls = 0
    cancelled = False

    def maybe_cancel_file_sync(path):
        nonlocal cancelled
        if failure_phase == "tree_sync" and not cancelled:
            cancelled = True
            raise SimulatedCancellation()
        return real_sync_file(path)

    def maybe_cancel_replace(source, target):
        nonlocal replace_calls, cancelled
        replace_calls += 1
        target_call = 1 if failure_phase == "backup_replace" else 2
        if failure_phase in {"backup_replace", "promotion_replace"}:
            if replace_calls == target_call and not cancelled:
                cancelled = True
                raise SimulatedCancellation()
        return real_replace(source, target)

    def maybe_cancel_directory_sync(path):
        nonlocal directory_sync_calls, cancelled
        directory_sync_calls += 1
        target_call = 2 if failure_phase == "backup_parent_sync" else 3
        if failure_phase in {"backup_parent_sync", "promotion_parent_sync"}:
            if directory_sync_calls == target_call and not cancelled:
                cancelled = True
                raise SimulatedCancellation()
        return real_sync_directory(path)

    monkeypatch.setattr(persistence, "_sync_file", maybe_cancel_file_sync)
    monkeypatch.setattr(persistence.os, "replace", maybe_cancel_replace)
    monkeypatch.setattr(persistence, "_sync_directory", maybe_cancel_directory_sync)

    with pytest.raises(SimulatedCancellation):
        with _atomic_checkpoint_directory(destination) as staging:
            (staging / "new.txt").write_text("first attempt", encoding="utf-8")

    assert sentinel.read_text(encoding="utf-8") == "preserve me"
    assert _temporary_artifacts(tmp_path, destination.name) == []

    with _atomic_checkpoint_directory(destination) as staging:
        (staging / "new.txt").write_text("retry", encoding="utf-8")

    assert not sentinel.exists()
    assert (destination / "new.txt").read_text(encoding="utf-8") == "retry"
    assert _temporary_artifacts(tmp_path, destination.name) == []


def test_atomic_checkpoint_restores_destination_on_promotion_failure(
    tmp_path,
    monkeypatch,
):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    sentinel = destination / "sentinel.txt"
    sentinel.write_text("preserve me", encoding="utf-8")
    real_replace = os.replace

    def fail_staging_promotion(source, target):
        if ".staging-" in Path(source).name:
            raise OSError("simulated promotion failure")
        return real_replace(source, target)

    monkeypatch.setattr(persistence.os, "replace", fail_staging_promotion)
    with pytest.raises(OSError, match="simulated promotion failure"):
        with _atomic_checkpoint_directory(destination) as staging:
            (staging / "new.txt").write_text("new", encoding="utf-8")

    assert sentinel.read_text() == "preserve me"
    assert _temporary_artifacts(tmp_path, destination.name) == []


def test_atomic_checkpoint_reports_recoverable_backup_when_rollback_fails(
    tmp_path,
    monkeypatch,
):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    (destination / "sentinel.txt").write_text("recover me", encoding="utf-8")
    real_replace = os.replace
    replacement_calls = 0

    def fail_promotion_and_restore(source, target):
        nonlocal replacement_calls
        replacement_calls += 1
        if replacement_calls >= 2:
            raise OSError("simulated replace failure")
        return real_replace(source, target)

    monkeypatch.setattr(persistence.os, "replace", fail_promotion_and_restore)
    with pytest.raises(RuntimeError, match="recover the previous checkpoint from"):
        with _atomic_checkpoint_directory(destination) as staging:
            (staging / "new.txt").write_text("new", encoding="utf-8")

    backups = list(tmp_path.glob(".checkpoint.backup-*"))
    assert len(backups) == 1
    assert (backups[0] / "sentinel.txt").read_text() == "recover me"


def test_atomic_checkpoint_reports_staging_cleanup_failure_without_masking_write(
    tmp_path,
    monkeypatch,
    caplog,
):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    sentinel = destination / "sentinel.txt"
    sentinel.write_text("preserve me", encoding="utf-8")
    real_remove = persistence._remove_checkpoint_path

    def fail_staging_cleanup(path):
        if ".staging-" in path.name:
            raise PermissionError("simulated cleanup denial")
        return real_remove(path)

    monkeypatch.setattr(persistence, "_remove_checkpoint_path", fail_staging_cleanup)
    with pytest.raises(OSError, match="simulated write failure"):
        with _atomic_checkpoint_directory(destination) as staging:
            (staging / "partial.bin").write_bytes(b"partial")
            raise OSError("simulated write failure")

    assert sentinel.read_text(encoding="utf-8") == "preserve me"
    assert caplog.messages == [
        "Checkpoint write failed, and staging directory "
        f"{staging} could not be removed: simulated cleanup denial",
    ]
    assert len(list(tmp_path.glob(".checkpoint.staging-*"))) == 1


def test_rebirth_failure_preserves_checkpoint_and_owned_offload(tmp_path):
    destination = tmp_path / "checkpoint"
    destination.mkdir()
    sentinel = destination / "sentinel.txt"
    sentinel.write_text("preserve me", encoding="utf-8")
    offload = tmp_path / "owned-offload"
    offload.mkdir()
    (offload / "weight.bin").write_bytes(b"still needed")

    pipeline = AbliterationPipeline(
        model_name="test-model",
        output_dir=str(destination),
        method="basic",
    )
    pipeline._on_log = lambda _message: None
    pipeline._on_stage = lambda _event: None
    pipeline.handle = MagicMock()
    pipeline.handle.model.state_dict.return_value = {"weight": torch.ones(1)}
    pipeline.handle.model.save_pretrained.side_effect = OSError("disk vanished")
    pipeline.handle._offload_dir = str(offload)
    pipeline.handle._owns_offload_dir = True

    with pytest.raises(OSError, match="disk vanished"):
        pipeline._rebirth()

    assert sentinel.read_text() == "preserve me"
    assert (offload / "weight.bin").read_bytes() == b"still needed"
    assert _temporary_artifacts(tmp_path, destination.name) == []


def test_cleanup_only_removes_pipeline_owned_offload_directory(tmp_path):
    caller_owned = tmp_path / "caller-owned"
    caller_owned.mkdir()
    (caller_owned / "weight.bin").write_bytes(b"owned by caller")
    pipeline = AbliterationPipeline(model_name="test-model", method="basic")
    pipeline._on_log = lambda _message: None
    pipeline.handle = MagicMock()
    pipeline.handle._offload_dir = str(caller_owned)
    pipeline.handle._owns_offload_dir = False

    pipeline._cleanup_offload_dir()

    assert (caller_owned / "weight.bin").read_bytes() == b"owned by caller"


def test_cleanup_removes_and_clears_pipeline_owned_offload_directory(tmp_path):
    owned = tmp_path / "pipeline-owned"
    owned.mkdir()
    (owned / "weight.bin").write_bytes(b"temporary")
    pipeline = AbliterationPipeline(model_name="test-model", method="basic")
    pipeline._on_log = lambda _message: None
    pipeline.handle = MagicMock()
    pipeline.handle._offload_dir = str(owned)
    pipeline.handle._owns_offload_dir = True

    pipeline._cleanup_offload_dir()

    assert not owned.exists()
    assert pipeline.handle._offload_dir is None
    assert pipeline.handle._owns_offload_dir is False


def test_cleanup_clears_stale_owned_offload_reference(tmp_path):
    pipeline = AbliterationPipeline(model_name="test-model", method="basic")
    pipeline._on_log = lambda _message: None
    pipeline.handle = MagicMock()
    pipeline.handle._offload_dir = str(tmp_path / "already-gone")
    pipeline.handle._owns_offload_dir = True

    pipeline._cleanup_offload_dir()

    assert pipeline.handle._offload_dir is None
    assert pipeline.handle._owns_offload_dir is False


def test_cleanup_failure_retains_owned_offload_reference_for_retry(
    tmp_path,
    monkeypatch,
):
    owned = tmp_path / "pipeline-owned"
    owned.mkdir()
    (owned / "weight.bin").write_bytes(b"temporary")
    messages = []
    pipeline = AbliterationPipeline(model_name="test-model", method="basic")
    pipeline._on_log = messages.append
    pipeline.handle = MagicMock()
    pipeline.handle._offload_dir = str(owned)
    pipeline.handle._owns_offload_dir = True

    def deny_cleanup(_path):
        raise PermissionError("simulated cleanup denial")

    monkeypatch.setattr(abliterate.shutil, "rmtree", deny_cleanup)

    pipeline._cleanup_offload_dir()

    assert (owned / "weight.bin").read_bytes() == b"temporary"
    assert pipeline.handle._offload_dir == str(owned)
    assert pipeline.handle._owns_offload_dir is True
    assert any("retaining the owned path for retry" in message for message in messages)
