from __future__ import annotations

import json
import os
from pathlib import Path
import pytest

from obliteratus.run_archive import RunArchive, _worker


class FakeProcess:
    def __init__(self, pid: int):
        self.pid = pid


def test_launch_persists_manifest_before_worker_and_redacts_secrets(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    observed = {}

    def fake_popen(command, **kwargs):
        run_id = command[command.index("--run-id") + 1]
        observed["manifest"] = archive._load(run_id)
        observed["command"] = command
        observed["kwargs"] = kwargs
        return FakeProcess(os.getpid())

    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 123)
    run_id = archive.launch(
        ["org/model", "--token", "private-value", "--dataset", "builtin"],
        notes="retain this experiment",
        popen=fake_popen,
    )

    assert observed["manifest"]["status"] == "queued"
    assert observed["manifest"]["phase"] == "manifested"
    assert observed["manifest"]["arguments"][2] == "[REDACTED]"
    assert observed["kwargs"]["start_new_session"] is True
    assert (tmp_path / run_id / "notes.md").read_text() == "retain this experiment\n"
    assert archive.status(run_id)["phase"] == "worker_started"


def test_launch_rejects_archive_escape_and_managed_output_override(tmp_path):
    archive = RunArchive(tmp_path)
    with pytest.raises(ValueError, match="model ID"):
        archive.launch(["--method", "advanced"])
    with pytest.raises(ValueError, match="managed"):
        archive.launch(["org/model", "--output-dir", "/tmp/escape"])
    with pytest.raises(ValueError, match="invalid run ID"):
        archive.status("../outside")


def test_archive_and_run_directory_symlinks_fail_closed(tmp_path):
    actual = tmp_path / "actual"
    actual.mkdir()
    alias = tmp_path / "alias"
    alias.symlink_to(actual, target_is_directory=True)
    with pytest.raises(ValueError, match="root must not be a symlink"):
        RunArchive(alias)

    archive = RunArchive(actual)
    run_id = "run-" + "c" * 32
    (actual / run_id).symlink_to(tmp_path, target_is_directory=True)
    with pytest.raises(ValueError, match="run directory must not be a symlink"):
        archive.status(run_id)


def test_status_recovers_nonterminal_run_when_worker_identity_is_lost(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 111)
    run_id = archive.launch(["org/model"], popen=lambda *_args, **_kwargs: FakeProcess(4242))
    monkeypatch.setattr(archive, "_worker_matches", lambda _manifest: False)

    status = archive.status(run_id)

    assert status["status"] == "failed"
    assert status["phase"] == "worker_lost"
    assert status["failure"]["type"] == "WorkerLost"


def test_cancel_validates_identity_then_signals_only_worker_group(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 222)
    run_id = archive.launch(["org/model"], popen=lambda *_args, **_kwargs: FakeProcess(5252))
    monkeypatch.setattr(archive, "_worker_matches", lambda _manifest: True)
    signals = []
    monkeypatch.setattr("obliteratus.run_archive.os.killpg", lambda pid, sig: signals.append((pid, sig)))

    status = archive.cancel(run_id)

    assert signals and signals[0][0] == 5252
    assert status["status"] == "cancelling"
    assert status["phase"] == "cancellation_requested"


def test_result_requires_terminal_state(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 333)
    run_id = archive.launch(["org/model"], popen=lambda *_args, **_kwargs: FakeProcess(6262))
    monkeypatch.setattr(archive, "_worker_matches", lambda _manifest: True)
    with pytest.raises(RuntimeError, match="not complete"):
        archive.result(run_id)


def test_pruning_preserves_notes_logs_manifest_failure_and_hashed_inventory(tmp_path):
    archive = RunArchive(tmp_path)
    run_id = "run-" + "a" * 32
    run_dir = tmp_path / run_id
    checkpoint = run_dir / "checkpoint"
    checkpoint.mkdir(parents=True)
    (checkpoint / "weights.bin").write_bytes(b"weights")
    (run_dir / "notes.md").write_text("failed candidate\n")
    (run_dir / "run.log").write_text("failure details\n")
    manifest = {
        "schema_version": 1,
        "run_id": run_id,
        "status": "failed",
        "phase": "verify",
        "failure": {"type": "QualityError", "message": "dominated", "phase": "verify"},
    }
    archive._save(manifest)

    result = archive.prune_checkpoint(run_id, reason="dominated run under disk pressure")

    assert not checkpoint.exists()
    assert (run_dir / "notes.md").is_file()
    assert (run_dir / "run.log").is_file()
    assert (run_dir / "manifest.json").is_file()
    inventory = json.loads((run_dir / "artifact-inventory.json").read_text())
    assert any(item["path"] == "checkpoint/weights.bin" for item in inventory["artifacts"])
    assert result["failure"]["message"] == "dominated"
    assert result["checkpoint_pruned"]["reason"].startswith("dominated")


def test_concurrent_run_ids_and_manifests_do_not_collide(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 444)
    counter = iter(range(7000, 7020))
    fake = lambda *_args, **_kwargs: FakeProcess(next(counter))

    run_ids = {archive.launch(["org/model"], popen=fake) for _ in range(20)}

    assert len(run_ids) == 20
    assert {archive._load(run_id)["run_id"] for run_id in run_ids} == run_ids


def test_checkpoint_metrics_are_returned_from_success_metadata(tmp_path):
    from obliteratus.run_archive import _checkpoint_metrics

    checkpoint = tmp_path / "checkpoint"
    checkpoint.mkdir()
    (checkpoint / "abliteration_metadata.json").write_text(
        json.dumps({"quality_metrics": {"refusal_rate": 0.21, "coherence": 0.8}})
    )
    assert _checkpoint_metrics(checkpoint) == {"refusal_rate": 0.21, "coherence": 0.8}


def test_list_skips_corrupt_manifests(tmp_path):
    archive = RunArchive(tmp_path)
    corrupt = tmp_path / ("run-" + "b" * 32)
    corrupt.mkdir()
    (corrupt / "manifest.json").write_text("[]")
    assert archive.list() == []


def test_default_root_honors_service_archive_environment(tmp_path, monkeypatch):
    monkeypatch.setenv("OBLITERATUS_RUN_ARCHIVE", str(tmp_path / "durable"))
    assert RunArchive().root == (tmp_path / "durable").resolve()


def test_dataset_manifest_records_hash_and_counts_without_prompt_text(tmp_path):
    archive = RunArchive(tmp_path)
    run_id = archive.begin(["org/model"])
    archive.record_dataset(
        run_id,
        identifier="builtin",
        harmful=["sensitive harmful prompt"],
        harmless=["ordinary harmless prompt"],
    )
    raw = (tmp_path / run_id / "manifest.json").read_text()
    manifest = json.loads(raw)
    assert "sensitive harmful prompt" not in raw
    assert manifest["dataset_inputs"][0]["harmful_count"] == 1
    assert len(manifest["dataset_inputs"][0]["sha256"]) == 64


def test_experiment_protocol_is_durable_without_raw_prompts(tmp_path):
    archive = RunArchive(tmp_path / "runs")
    run_id = archive.begin(["model"])
    protocol = {
        "protocol": "qwen38-v1",
        "manifest_sha256": "a" * 64,
        "counts": {"train": 500, "tune": 142, "test": 200},
        "pair_ids": {"train": ["b" * 64], "tune": [], "test": []},
    }

    manifest = archive.record_experiment_protocol(run_id, protocol)

    assert manifest["experiment_protocol"]["protocol"] == "qwen38-v1"
    path = tmp_path / "runs" / run_id / "experiment-protocol.json"
    assert json.loads(path.read_text(encoding="utf-8")) == protocol


def test_failure_detail_redacts_huggingface_and_bearer_tokens(tmp_path):
    archive = RunArchive(tmp_path)
    run_id = archive.begin(["org/model"])
    result = archive.fail(
        run_id,
        RuntimeError("request used hf_abcdefghijklmnopqrstuvwxyz and Bearer abcdefghijklmnop"),
        phase="download",
    )
    assert result["failure"]["message"].count("[REDACTED]") == 2
    assert "hf_" not in result["failure"]["message"]


class FakeLifecycle:
    def __init__(self):
        self.events = []

    def loading(self, model):
        self.events.append(("loading", model))

    def ready(self):
        self.events.append(("ready", None))

    def resize(self, memory):
        self.events.append(("resize", memory.reserved_bytes))

    def release(self, *, reason):
        self.events.append(("release", reason))


def _embedded_run(archive):
    return archive.begin(["org/model"], metadata={"configuration": {"seed": 7}})


def test_worker_success_writes_metrics_inventory_and_atomic_marker(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    run_id = _embedded_run(archive)
    lifecycle = FakeLifecycle()
    monkeypatch.setattr("obliteratus.run_archive.from_environment", lambda: lifecycle)

    class SuccessfulChild:
        def __init__(self, command, **_kwargs):
            checkpoint = Path(command[command.index("--output-dir") + 1])
            checkpoint.mkdir()
            (checkpoint / "weights.bin").write_bytes(b"model")
            (checkpoint / "abliteration_metadata.json").write_text(
                json.dumps({"quality_metrics": {"refusal_rate": 0.2, "coherence": 0.8}})
            )

        def wait(self, timeout=None):
            return 0

        def poll(self):
            return 0

    monkeypatch.setattr("obliteratus.run_archive.subprocess.Popen", SuccessfulChild)

    assert _worker(archive, run_id, ["org/model"]) == 0
    result = archive.result(run_id)
    assert result["result"]["metrics"] == {"refusal_rate": 0.2, "coherence": 0.8}
    assert (tmp_path / run_id / "COMPLETE").is_file()
    assert lifecycle.events == [
        ("loading", "org/model"), ("resize", 0), ("ready", None),
        ("release", "complete")
    ]


def test_worker_failure_preserves_partial_checkpoint_and_failure_detail(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    run_id = _embedded_run(archive)
    lifecycle = FakeLifecycle()
    monkeypatch.setattr("obliteratus.run_archive.from_environment", lambda: lifecycle)

    class FailedChild:
        def __init__(self, command, **_kwargs):
            checkpoint = Path(command[command.index("--output-dir") + 1])
            checkpoint.mkdir()
            (checkpoint / "partial.bin").write_bytes(b"partial")

        def wait(self, timeout=None):
            return 17

        def poll(self):
            return 17

    monkeypatch.setattr("obliteratus.run_archive.subprocess.Popen", FailedChild)

    assert _worker(archive, run_id, ["org/model"]) == 1
    result = archive.result(run_id)
    assert result["status"] == "failed"
    assert result["failure"]["type"] == "RuntimeError"
    assert "status 17" in result["failure"]["message"]
    assert (tmp_path / run_id / "checkpoint" / "partial.bin").is_file()
    assert lifecycle.events[-1] == ("release", "failed")


def test_worker_commits_cancelled_state_after_cooperative_interrupt(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    run_id = _embedded_run(archive)
    lifecycle = FakeLifecycle()
    monkeypatch.setattr("obliteratus.run_archive.from_environment", lambda: lifecycle)

    class CancelledChild:
        def __init__(self, _command, **_kwargs):
            self.terminated = False

        def wait(self, timeout=None):
            if timeout is not None:
                return -15
            manifest = archive._load(run_id)
            manifest["status"] = "cancelling"
            archive._save(manifest)
            raise KeyboardInterrupt("cancelled")

        def poll(self):
            return None if not self.terminated else -15

        def terminate(self):
            self.terminated = True

    monkeypatch.setattr("obliteratus.run_archive.subprocess.Popen", CancelledChild)

    assert _worker(archive, run_id, ["org/model"]) == 130
    result = archive.result(run_id)
    assert result["status"] == "cancelled"
    assert result["phase"] == "cancelled"
    assert lifecycle.events[-1] == ("release", "cancelled")


def test_restart_recovery_preserves_partial_save_and_logs(tmp_path, monkeypatch):
    archive = RunArchive(tmp_path)
    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: 555)
    run_id = archive.launch(["org/model"], popen=lambda *_args, **_kwargs: FakeProcess(9999))
    run_dir = tmp_path / run_id
    (run_dir / "checkpoint").mkdir()
    (run_dir / "checkpoint" / "partial.bin").write_bytes(b"partial")
    (run_dir / "run.log").write_text("last durable phase\n")
    monkeypatch.setattr("obliteratus.run_archive._process_start_ticks", lambda _pid: None)

    recovered = RunArchive(tmp_path).status(run_id)

    assert recovered["status"] == "failed"
    assert recovered["failure"]["type"] == "WorkerLost"
    assert (run_dir / "checkpoint" / "partial.bin").is_file()
    assert (run_dir / "run.log").read_text() == "last durable phase\n"