"""Tests for first-class control-plane lease lifecycle (#601)."""

from __future__ import annotations

import os
import tempfile
import unittest
from datetime import timedelta
from unittest import mock

from allocator_service import allocate_next_work, WorkCandidate
from control_plane_db import ControlPlaneDB, ForeignLeaseError, _ts, _utc_now
import lease_lifecycle as ll
import merger_lease_adoption as mla
import reviewer_pr_lease as rpl
from datetime import datetime, timezone


class LeaseLifecycleTest(unittest.TestCase):
    def setUp(self) -> None:
        self._tmp = tempfile.TemporaryDirectory()
        self.db_path = os.path.join(self._tmp.name, "cp.sqlite3")
        self.db = ControlPlaneDB(self.db_path)
        self.db.upsert_session(session_id="owner", role="author", profile="prgs-author", pid=111)
        self.db.upsert_session(session_id="other", role="author", profile="prgs-author", pid=222)

    def tearDown(self) -> None:
        self._tmp.cleanup()

    def _assign(self, session_id="owner", number=601, **kwargs):
        return self.db.assign_and_lease(
            session_id=session_id,
            role="author",
            remote="prgs",
            org="Scaled-Tech-Consulting",
            repo="Gitea-Tools",
            kind="issue",
            number=number,
            allowed_actions=("implement", "comment", "push", "create_pr"),
            forbidden_actions=("approve", "merge", "self_select_without_assignment"),
            worktree_path=kwargs.pop("worktree_path", self._tmp.name),
            owner_pid=kwargs.pop("owner_pid", os.getpid()),
            **kwargs,
        )

    def test_list_active_leases_as_workflow_state(self) -> None:
        a = self._assign(number=601)
        self._assign(session_id="other", number=602)
        listed = ll.list_active_leases(
            self.db, remote="prgs", org="Scaled-Tech-Consulting", repo="Gitea-Tools"
        )
        self.assertTrue(listed["success"])
        self.assertEqual(listed["count"], 2)
        self.assertEqual(listed["authoritative_source"], "control_plane_db")
        self.assertFalse(listed["file_lock_only"])
        self.assertFalse(listed["comment_lease_only"])
        numbers = sorted(x["work_number"] for x in listed["leases"])
        self.assertEqual(numbers, [601, 602])
        self.assertIsNotNone(a.lease_id)

    def test_adopt_lease_owner_resume_preserves_provenance(self) -> None:
        a = self._assign()
        res = ll.adopt_lease(
            self.db,
            lease_id=a.lease_id,
            adopter_session_id="owner",
            role="author",
            worktree_path=self._tmp.name,
        )
        self.assertTrue(res["success"])
        self.assertTrue(res["same_owner"])
        prov = res["provenance"]
        self.assertEqual(prov["adopted_from_session_id"], "owner")
        self.assertEqual(prov["adopted_by_session_id"], "owner")
        self.assertEqual(prov["work_number"], 601)
        self.assertEqual(prov["worktree_path"], self._tmp.name)
        state = self.db.get_lease_workflow_state(a.lease_id)
        self.assertIsNotNone(state)
        self.assertEqual(state["lease"]["status"], "active")

    def test_release_lease_explicit_recorded(self) -> None:
        a = self._assign()
        res = ll.release_lease(self.db, lease_id=a.lease_id, session_id="owner")
        self.assertEqual(res["outcome"], "released")
        self.assertIn("release_proof", res)
        self.assertEqual(res["release_proof"]["status"], "released")
        state = self.db.get_lease_workflow_state(a.lease_id)
        self.assertEqual(state["lease"]["status"], "released")

    def test_expire_and_reclaim_expired_lease(self) -> None:
        a = self._assign(lease_ttl_seconds=1)
        past = _utc_now() - timedelta(hours=2)
        import sqlite3

        conn = sqlite3.connect(self.db_path)
        try:
            conn.execute(
                "UPDATE leases SET expires_at = ? WHERE lease_id = ?",
                (_ts(past), a.lease_id),
            )
            conn.commit()
        finally:
            conn.close()
        exp = ll.expire_leases(self.db)
        self.assertGreaterEqual(exp["expired_count"], 1)
        reclaimed = ll.reclaim_expired_lease(
            self.db,
            lease_id=a.lease_id,
            session_id="other",
            role="author",
            worktree_path=self._tmp.name,
        )
        self.assertEqual(reclaimed["outcome"], "reclaimed")
        self.assertEqual(reclaimed["assignment"]["session_id"], "other")
        self.assertEqual(
            reclaimed["provenance"]["adopted_from_session_id"], "owner"
        )
        self.assertEqual(
            reclaimed["provenance"]["adopted_by_session_id"], "other"
        )

    def test_abandon_stale_lease_with_required_proof(self) -> None:
        a = self._assign(owner_pid=99999999, worktree_path="/nonexistent/path/for-601")
        # Force active but dead pid + missing worktree
        proof = ll.AbandonProof(
            dead_process=True,
            missing_worktree=True,
            no_open_pr=True,
            no_live_mutation_risk=True,
            owner_pid=99999999,
            worktree_path="/nonexistent/path/for-601",
        )
        res = ll.abandon_lease(
            self.db,
            lease_id=a.lease_id,
            requester_session_id="other",
            proof=proof,
        )
        self.assertEqual(res["outcome"], "abandoned")
        self.assertEqual(res["prior_owner_session_id"], "owner")
        self.assertTrue(res["abandon_proof"]["dead_process"])
        state = self.db.get_lease_workflow_state(a.lease_id)
        self.assertEqual(state["lease"]["status"], "abandoned")

    def test_refuse_steal_active_foreign_lease(self) -> None:
        a = self._assign()
        with self.assertRaises(ll.LeaseLifecycleError) as ctx:
            ll.adopt_lease(
                self.db,
                lease_id=a.lease_id,
                adopter_session_id="other",
                role="author",
            )
        self.assertIn("steal", str(ctx.exception).lower())
        decision = ll.inspect_lease(
            self.db, a.lease_id, caller_session_id="other"
        )
        self.assertEqual(decision["safe_next_action"], ll.SAFE_WAIT_FOREIGN)
        self.assertTrue(decision["block"])

    def test_refuse_ambiguous_lease_ownership_stale_id(self) -> None:
        decision = ll.inspect_lease(
            self.db, "lease-does-not-exist", caller_session_id="owner"
        )
        self.assertFalse(decision["found"])
        self.assertEqual(decision["safe_next_action"], ll.SAFE_STALE_PROMPT)
        with self.assertRaises(ll.LeaseLifecycleError):
            ll.adopt_lease(
                self.db,
                lease_id="lease-does-not-exist",
                adopter_session_id="owner",
                role="author",
            )

    def test_provenance_on_adopt_release_abandon(self) -> None:
        a = self._assign()
        adopted = ll.adopt_lease(
            self.db,
            lease_id=a.lease_id,
            adopter_session_id="owner",
            role="author",
            worktree_path=self._tmp.name,
            expected_head_sha="abc",
        )
        self.assertIn("adopted_from_session_id", adopted["provenance"])
        self.assertIn("adopted_by_session_id", adopted["provenance"])
        released = ll.release_lease(
            self.db, lease_id=a.lease_id, session_id="owner"
        )
        self.assertEqual(released["release_proof"]["session_id"], "owner")

        b = self._assign(number=700, owner_pid=1, worktree_path="/no/such/wt")
        abandoned = ll.abandon_lease(
            self.db,
            lease_id=b.lease_id,
            requester_session_id="owner",
            proof=ll.AbandonProof(
                dead_process=True,
                missing_worktree=True,
                no_live_mutation_risk=True,
                no_open_pr=True,
            ),
        )
        self.assertIn("abandon_proof", abandoned)
        self.assertEqual(abandoned["abandon_proof"]["dead_process"], True)

    def test_allocator_assignment_lease_compatible(self) -> None:
        """Allocator assign+lease still works and is listable/inspectable."""
        cands = [
            WorkCandidate(
                kind="issue",
                number=601,
                labels=("status:ready",),
                title="leases",
                priority=20,
            )
        ]
        res = allocate_next_work(
            self.db,
            session_id="alloc-s",
            role="author",
            remote="prgs",
            org="Scaled-Tech-Consulting",
            repo="Gitea-Tools",
            candidates=cands,
            apply=True,
            profile_name="prgs-author",
            username="jcwalker3",
        )
        self.assertEqual(res["outcome"], "assigned_work")
        lid = res["assignment"]["lease_id"]
        listed = ll.list_active_leases(self.db, remote="prgs")
        ids = [x["lease_id"] for x in listed["leases"]]
        self.assertIn(lid, ids)
        insp = ll.inspect_lease(self.db, lid, caller_session_id="alloc-s")
        self.assertTrue(insp["found"])
        self.assertIn(
            insp["safe_next_action"],
            (ll.SAFE_OWNER_RESUME, ll.SAFE_ABANDON_ALLOWED),
        )

    def test_reviewer_merger_lease_handoff_still_works(self) -> None:
        """#536 merger adoption path is independent and must not regress."""
        rpl.clear_session_lease()
        body = mla.format_adoption_body(
            repo="Scaled-Tech-Consulting/Gitea-Tools",
            pr_number=999,
            issue_number=601,
            adopter_identity="sysadmin",
            adopter_profile="prgs-merger",
            adopter_session_id="merger-1",
            worktree="branches/merge-pr999",
            candidate_head="a" * 40,
            target_branch="master",
            target_branch_sha="b" * 40,
            adopted_from_session_id="reviewer-1",
            adopted_from_profile="prgs-reviewer",
            adopted_from_reviewer_identity="sysadmin",
            adopted_from_comment_id=42,
        )
        self.assertIn(mla.ADOPTION_MARKER, body)
        self.assertIn("adopted_from_session_id: reviewer-1", body)
        self.assertIn("adopted_by_profile: prgs-merger", body)
        prov = mla.build_lease_provenance(
            source=mla.SOURCE_ADOPT,
            comment_id=42,
            adopted_from_session_id="reviewer-1",
            adoption_reason=mla.DEFAULT_ADOPTION_REASON,
        )
        rpl.record_session_lease(
            {
                "pr_number": 999,
                "session_id": "merger-1",
                "comment_id": 42,
            },
            lease_provenance=prov,
        )
        proof = mla.describe_session_lease_proof(rpl.get_session_lease())
        self.assertTrue(proof["lease_proof_sanctioned"])
        self.assertEqual(proof["lease_proof_source"], mla.SOURCE_ADOPT)
        rpl.clear_session_lease()

    def test_missing_worktree_and_dead_pid_handled_safely(self) -> None:
        a = self._assign(owner_pid=1, worktree_path="/definitely/missing/wt-601")
        with mock.patch.object(ll, "is_process_alive", return_value=False):
            fr = ll.classify_lease_freshness(
                self.db.get_lease_workflow_state(a.lease_id)["lease"]
            )
        self.assertIn(fr["freshness"], ("stale_dead_process", "stale_missing_worktree"))
        # Insufficient abandon proof fails closed
        with self.assertRaises(ll.LeaseLifecycleError):
            ll.abandon_lease(
                self.db,
                lease_id=a.lease_id,
                requester_session_id="other",
                proof=ll.AbandonProof(dead_process=True),  # missing no_live_mutation_risk
            )

    def test_file_lock_or_comment_not_authoritative_alone(self) -> None:
        report = ll.non_db_lease_authority_report(
            file_lock_present=True,
            comment_lease_present=False,
            db_lease_present=False,
        )
        self.assertEqual(report["safe_next_action"], ll.SAFE_NO_AUTHORITY)
        self.assertTrue(report["file_lock_only"])
        self.assertIsNone(report["authoritative_source"])
        report2 = ll.non_db_lease_authority_report(
            file_lock_present=True,
            comment_lease_present=True,
            db_lease_present=True,
        )
        self.assertEqual(report2["authoritative_source"], "control_plane_db")
        self.assertFalse(report2["file_lock_only"])
        self.assertFalse(report2["comment_lease_only"])

    def test_abandon_proof_is_sufficient_rules(self) -> None:
        self.assertFalse(ll.AbandonProof(dead_process=True).is_sufficient())
        self.assertTrue(
            ll.AbandonProof(
                dead_process=True,
                missing_worktree=True,
                no_live_mutation_risk=True,
                no_open_pr=True,
            ).is_sufficient()
        )
        self.assertTrue(
            ll.AbandonProof(
                dead_process=True,
                no_live_mutation_risk=True,
                same_owner=True,
            ).is_sufficient()
        )
        self.assertTrue(
            ll.AbandonProof(
                missing_worktree=True,
                no_live_mutation_risk=True,
                operator_authorized=True,
            ).is_sufficient()
        )


class IssueLockExpiredReclaimTest(unittest.TestCase):
    def test_expired_lock_reclaim_allowed_when_dead_and_missing_wt(self) -> None:
        import issue_lock_store as ils

        lock = {
            "issue_number": 601,
            "branch_name": "feat/issue-601-x",
            "worktree_path": "/no/such/worktree-601",
            "session_pid": 1,
            "work_lease": {
                "operation_type": "author_issue_work",
                "expires_at": "2000-01-01T00:00:00Z",
            },
        }
        with mock.patch.object(ils, "is_process_alive", return_value=False):
            res = ils.assess_expired_lock_reclaim(lock)
        self.assertTrue(res["reclaim_allowed"])
        # Conflict assessor allows takeover
        block = ils.assess_same_issue_lease_conflict(
            lock,
            issue_number=601,
            branch_name="feat/issue-601-first-class-leases",
            worktree_path="/tmp/new",
        )
        self.assertIsNone(block)

    def test_live_lock_reclaim_refused(self) -> None:
        import issue_lock_store as ils
        from datetime import datetime, timedelta, timezone

        future = (datetime.now(timezone.utc) + timedelta(hours=2)).strftime(
            "%Y-%m-%dT%H:%M:%SZ"
        )
        lock = {
            "issue_number": 601,
            "branch_name": "feat/issue-601-x",
            "worktree_path": tempfile.gettempdir(),
            "session_pid": os.getpid(),
            "work_lease": {
                "operation_type": "author_issue_work",
                "expires_at": future,
                "last_heartbeat_at": future,
            },
        }
        res = ils.assess_expired_lock_reclaim(lock)
        self.assertFalse(res["reclaim_allowed"])


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