"""Tests for the pre-restart drain proof and hard gate (#661).

Covers the acceptance criteria:

1. Restart apply without a proof fails closed.
2. A successful drain produces a verifiable proof.
3. An open unsafe mutation makes the proof fail (multi-session fixture).
4. Pass / fail / expired verification paths.

Plus the security posture: forged/tampered proofs are rejected, break-glass is
the only bypass and is never silent, a stale blast-radius fingerprint rejects a
proof, and no per-process secret ever leaks into a serialized artifact.
"""

from __future__ import annotations

import os
import unittest
from datetime import datetime, timedelta, timezone

import drain_proof as dp
import restart_coordinator as rc


NOW = datetime(2026, 7, 24, 6, 0, 0, tzinfo=timezone.utc)
SECRET = b"unit-test-drain-proof-secret-0123456789abcdef"


def _live_pid() -> int:
    return os.getpid()


def _clean_drain_state() -> dict:
    """Every drain action succeeded, no sessions outstanding."""

    return {
        "assignments_stopped": True,
        "checkpoints_complete": True,
        "handoffs_verified": True,
        "leases_handled": True,
        "acks": {},  # no other live sessions to acknowledge
        "ack_timeout_policy_applied": False,
    }


def _safe_report() -> dict:
    """Impact report with no other live work: a restart here is safe."""

    report = rc.evaluate_restart_impact(
        {"sessions": [], "leases": [], "inventory_complete": True},
        now=NOW,
        requesting_session_id="prgs-controller-1-req",
    )
    return report.as_dict()


def _unsafe_mutation_report() -> dict:
    """Multi-session report: a second session holds a live author mutation."""

    sessions = [
        {
            "session_id": "prgs-controller-1-req",
            "role": "controller",
            "profile": "prgs-controller",
            "pid": _live_pid(),
            "status": "active",
            "last_heartbeat_at": NOW.isoformat(),
        },
        {
            "session_id": "prgs-author-99",
            "role": "author",
            "profile": "prgs-author",
            "pid": _live_pid(),
            "status": "active",
            "last_heartbeat_at": NOW.isoformat(),
        },
    ]
    leases = [
        {
            "lease_id": "lease-mut",
            "session_id": "prgs-author-99",
            "role": "author",
            "phase": "implementing",
            "work_kind": "issue",
            "work_number": 661,
            "worktree_path": "branches/issue-661",
            "freshness": {"freshness": "active"},
        }
    ]
    report = rc.evaluate_restart_impact(
        {"sessions": sessions, "leases": leases, "inventory_complete": True},
        now=NOW,
        requesting_session_id="prgs-controller-1-req",
    )
    return report.as_dict()


class BuildDrainProofTests(unittest.TestCase):
    def test_clean_drain_produces_verifiable_clean_proof(self):
        """AC#2: a successful drain produces a verifiable proof."""

        proof = dp.build_drain_proof(
            impact_report=_safe_report(),
            drain_state=_clean_drain_state(),
            requesting_session_id="prgs-controller-1-req",
            now=NOW,
            secret=SECRET,
        )
        self.assertTrue(proof.clean)
        self.assertEqual(proof.failed_checks, [])
        self.assertEqual(
            {c.name for c in proof.checks}, set(dp.REQUIRED_CHECKS)
        )
        result = dp.verify_drain_proof(
            proof.as_dict(), now=NOW, secret=SECRET
        )
        self.assertTrue(result.valid, result.reasons)
        self.assertFalse(result.expired)
        self.assertFalse(result.tampered)

    def test_open_mutation_makes_proof_unclean(self):
        """AC#3: an unsafe mutation still in flight fails the proof."""

        proof = dp.build_drain_proof(
            impact_report=_unsafe_mutation_report(),
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        )
        self.assertFalse(proof.clean)
        self.assertIn(dp.CHECK_NO_INFLIGHT_MUTATIONS, proof.failed_checks)
        # Leases-handled also fails: the report still shows a disruptive lease.
        self.assertIn(dp.CHECK_LEASES_HANDLED, proof.failed_checks)
        result = dp.verify_drain_proof(proof.as_dict(), now=NOW, secret=SECRET)
        self.assertFalse(result.valid)

    def test_incomplete_inventory_fails_no_mutations_check(self):
        proof = dp.build_drain_proof(
            impact_report={"inventory_complete": False},
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        )
        self.assertFalse(proof.clean)
        self.assertIn(dp.CHECK_NO_INFLIGHT_MUTATIONS, proof.failed_checks)

    def test_missing_checkpoint_flag_fails_closed(self):
        state = _clean_drain_state()
        del state["checkpoints_complete"]
        proof = dp.build_drain_proof(
            impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET
        )
        self.assertFalse(proof.clean)
        self.assertIn(dp.CHECK_CHECKPOINTS_COMPLETE, proof.failed_checks)

    def test_non_true_flags_fail_closed(self):
        """A truthy-but-not-True value (e.g. the string 'yes') must not pass."""

        state = _clean_drain_state()
        state["assignments_stopped"] = "yes"
        proof = dp.build_drain_proof(
            impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET
        )
        self.assertIn(dp.CHECK_ASSIGNMENTS_STOPPED, proof.failed_checks)

    def test_ack_timeout_policy_satisfies_ack_check(self):
        state = _clean_drain_state()
        state["acks"] = {"prgs-author-99": "pending"}
        state["ack_timeout_policy_applied"] = True
        proof = dp.build_drain_proof(
            impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET
        )
        names = {c.name: c.passed for c in proof.checks}
        self.assertTrue(names[dp.CHECK_ACKS_OR_TIMEOUT])

    def test_outstanding_acks_without_timeout_fail(self):
        state = _clean_drain_state()
        state["acks"] = {"prgs-author-99": "pending"}
        state["ack_timeout_policy_applied"] = False
        proof = dp.build_drain_proof(
            impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET
        )
        self.assertIn(dp.CHECK_ACKS_OR_TIMEOUT, proof.failed_checks)

    def test_all_acked_satisfies_ack_check(self):
        state = _clean_drain_state()
        state["acks"] = {"prgs-author-99": "acked", "prgs-author-2": "acknowledged"}
        proof = dp.build_drain_proof(
            impact_report=_safe_report(), drain_state=state, now=NOW, secret=SECRET
        )
        names = {c.name: c.passed for c in proof.checks}
        self.assertTrue(names[dp.CHECK_ACKS_OR_TIMEOUT])


class VerifyDrainProofTests(unittest.TestCase):
    def _clean_proof_dict(self) -> dict:
        return dp.build_drain_proof(
            impact_report=_safe_report(),
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        ).as_dict()

    def test_missing_proof_is_invalid(self):
        result = dp.verify_drain_proof(None, now=NOW, secret=SECRET)
        self.assertFalse(result.valid)
        self.assertIsNone(result.proof_id)

    def test_expired_proof_is_invalid(self):
        """AC#4: an expired proof fails verification."""

        proof = self._clean_proof_dict()
        later = NOW + timedelta(seconds=dp.DEFAULT_PROOF_TTL_SECONDS + 1)
        result = dp.verify_drain_proof(proof, now=later, secret=SECRET)
        self.assertFalse(result.valid)
        self.assertTrue(result.expired)

    def test_proof_valid_just_before_expiry(self):
        proof = self._clean_proof_dict()
        almost = NOW + timedelta(seconds=dp.DEFAULT_PROOF_TTL_SECONDS - 1)
        result = dp.verify_drain_proof(proof, now=almost, secret=SECRET)
        self.assertTrue(result.valid, result.reasons)

    def test_wrong_secret_rejected(self):
        """A proof minted in a prior process (different secret) will not verify."""

        proof = self._clean_proof_dict()
        result = dp.verify_drain_proof(proof, now=NOW, secret=b"other-secret")
        self.assertFalse(result.valid)
        self.assertTrue(result.tampered)

    def test_flipping_clean_flag_is_detected(self):
        """Forging clean=True on an unclean proof breaks the signature."""

        unclean = dp.build_drain_proof(
            impact_report=_unsafe_mutation_report(),
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        ).as_dict()
        self.assertFalse(unclean["clean"])
        unclean["clean"] = True  # forge
        result = dp.verify_drain_proof(unclean, now=NOW, secret=SECRET)
        self.assertFalse(result.valid)
        self.assertTrue(result.tampered)

    def test_tampering_a_check_is_detected(self):
        unclean = dp.build_drain_proof(
            impact_report=_unsafe_mutation_report(),
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        ).as_dict()
        for c in unclean["checks"]:
            if c["name"] == dp.CHECK_NO_INFLIGHT_MUTATIONS:
                c["passed"] = True  # forge the failing check to pass
        result = dp.verify_drain_proof(unclean, now=NOW, secret=SECRET)
        self.assertFalse(result.valid)
        self.assertTrue(result.tampered)

    def test_missing_required_check_rejected(self):
        proof = self._clean_proof_dict()
        proof["checks"] = [
            c for c in proof["checks"] if c["name"] != dp.CHECK_HANDOFFS_OK
        ]
        result = dp.verify_drain_proof(proof, now=NOW, secret=SECRET)
        self.assertFalse(result.valid)

    def test_stale_fingerprint_rejected(self):
        proof = self._clean_proof_dict()
        result = dp.verify_drain_proof(
            proof,
            now=NOW,
            secret=SECRET,
            expected_impact_fingerprint="deadbeef",
        )
        self.assertFalse(result.valid)

    def test_matching_fingerprint_accepted(self):
        report = _safe_report()
        proof = dp.build_drain_proof(
            impact_report=report,
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        ).as_dict()
        fp = dp.impact_fingerprint(report)
        result = dp.verify_drain_proof(
            proof, now=NOW, secret=SECRET, expected_impact_fingerprint=fp
        )
        self.assertTrue(result.valid, result.reasons)


class GateApplyRestartTests(unittest.TestCase):
    def _clean_proof_dict(self) -> dict:
        return dp.build_drain_proof(
            impact_report=_safe_report(),
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        ).as_dict()

    def test_apply_without_proof_denied(self):
        """AC#1: restart apply without a proof fails closed + raises incident."""

        decision = dp.gate_apply_restart(proof=None, now=NOW, secret=SECRET)
        self.assertFalse(decision.allow)
        self.assertEqual(decision.verdict, dp.GATE_DENY)
        self.assertIsNotNone(decision.incident)
        self.assertEqual(
            decision.incident["kind"], "restart_drain_gate_denied"
        )

    def test_apply_with_valid_proof_allowed(self):
        decision = dp.gate_apply_restart(
            proof=self._clean_proof_dict(), now=NOW, secret=SECRET
        )
        self.assertTrue(decision.allow)
        self.assertEqual(decision.verdict, dp.GATE_ALLOW)
        self.assertIsNone(decision.incident)

    def test_apply_with_expired_proof_denied_with_incident(self):
        later = NOW + timedelta(seconds=dp.DEFAULT_PROOF_TTL_SECONDS + 5)
        decision = dp.gate_apply_restart(
            proof=self._clean_proof_dict(), now=later, secret=SECRET
        )
        self.assertFalse(decision.allow)
        self.assertIsNotNone(decision.incident)

    def test_apply_with_unclean_proof_denied(self):
        """AC#3 at the gate: an unsafe-mutation proof is denied."""

        unclean = dp.build_drain_proof(
            impact_report=_unsafe_mutation_report(),
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        ).as_dict()
        decision = dp.gate_apply_restart(proof=unclean, now=NOW, secret=SECRET)
        self.assertFalse(decision.allow)
        self.assertIsNotNone(decision.incident)

    def test_break_glass_allows_without_proof_but_records_bypass(self):
        decision = dp.gate_apply_restart(
            proof=None, now=NOW, secret=SECRET, break_glass=True
        )
        self.assertTrue(decision.allow)
        self.assertEqual(decision.verdict, dp.GATE_BREAK_GLASS)
        self.assertTrue(decision.break_glass)
        self.assertIsNone(decision.incident)
        self.assertTrue(decision.audit_record["break_glass"])

    def test_denied_gate_carries_stale_fingerprint_reason(self):
        decision = dp.gate_apply_restart(
            proof=self._clean_proof_dict(),
            now=NOW,
            secret=SECRET,
            expected_impact_fingerprint="not-the-fingerprint",
        )
        self.assertFalse(decision.allow)


class SecretHygieneTests(unittest.TestCase):
    def test_secret_never_serialized(self):
        proof = dp.build_drain_proof(
            impact_report=_safe_report(),
            drain_state=_clean_drain_state(),
            now=NOW,
            secret=SECRET,
        )
        blob = dp._canonical(proof.as_dict())
        self.assertNotIn(SECRET.decode(), blob)
        # The signature is a hex digest, not the raw secret.
        self.assertNotIn(SECRET.hex(), blob)

    def test_incident_descriptor_has_no_secret(self):
        decision = dp.gate_apply_restart(proof=None, now=NOW, secret=SECRET)
        blob = dp._canonical(decision.incident)
        self.assertNotIn(SECRET.decode(), blob)


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