"""Tests for controller-owned allocator (#600) on control-plane DB (#613)."""

from __future__ import annotations

import os
import tempfile
import threading
import unittest
from concurrent.futures import ThreadPoolExecutor, as_completed
from datetime import datetime, timezone

from allocator_service import (
    OUTCOME_ASSIGNED,
    OUTCOME_BLOCKED_TERMINAL,
    OUTCOME_NO_SAFE,
    OUTCOME_PREVIEW,
    OUTCOME_WAIT,
    WorkCandidate,
    _drop_expired_claims,
    allocate_next_work,
    candidate_from_dict,
    classify_skip,
    expected_role_for_candidate,
)
from control_plane_db import ControlPlaneDB, InvalidWorkKindError


class AllocatorServiceTest(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)

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

    def _alloc(self, **kwargs):
        defaults = dict(
            db=self.db,
            session_id="s-test",
            role="author",
            remote="prgs",
            org="org",
            repo="repo",
            candidates=[],
            apply=False,
            profile_name="prgs-author",
            username="jcwalker3",
        )
        defaults.update(kwargs)
        return allocate_next_work(**defaults)

    def test_selects_ready_issue_for_author(self) -> None:
        cands = [
            WorkCandidate(
                kind="issue",
                number=612,
                labels=("status:ready",),
                title="bridge",
                dependency_unmet=True,
                dependency_reason="downstream of #600",
            ),
            WorkCandidate(
                kind="issue",
                number=600,
                labels=("status:ready", "type:feature"),
                title="allocator",
                priority=50,
            ),
            WorkCandidate(
                kind="issue",
                number=601,
                labels=("status:blocked",),
                title="blocked",
                blocked=True,
            ),
        ]
        res = self._alloc(candidates=cands, apply=False)
        self.assertTrue(res["success"])
        self.assertEqual(res["outcome"], OUTCOME_PREVIEW)
        self.assertEqual(res["selected"]["number"], 600)
        self.assertEqual(res["selected"]["kind"], "issue")
        # When 600 is highest priority valid work, lower-priority blocked/dep
        # candidates are not visited. Prove they are skipped when they sort first.
        res2 = self._alloc(
            candidates=[
                WorkCandidate(
                    kind="issue",
                    number=612,
                    labels=("status:ready",),
                    priority=99,
                    dependency_unmet=True,
                    dependency_reason="downstream of #600",
                ),
                WorkCandidate(
                    kind="issue",
                    number=601,
                    labels=("status:blocked",),
                    priority=98,
                    blocked=True,
                ),
                WorkCandidate(
                    kind="issue",
                    number=600,
                    labels=("status:ready",),
                    priority=1,
                ),
            ],
            apply=False,
        )
        skipped_nums = {s["number"] for s in res2["skipped"]}
        self.assertIn(612, skipped_nums)
        self.assertIn(601, skipped_nums)
        self.assertEqual(res2["selected"]["number"], 600)
        self.assertIsNone(res["assignment"])
        self.assertFalse(res["file_lock_only"])
        self.assertFalse(res["comment_lease_only"])

    def test_atomic_assign_and_lease_on_apply(self) -> None:
        cands = [
            WorkCandidate(
                kind="issue",
                number=600,
                labels=("status:ready",),
                priority=10,
            )
        ]
        res = self._alloc(candidates=cands, apply=True, session_id="s-a")
        self.assertEqual(res["outcome"], OUTCOME_ASSIGNED)
        asn = res["assignment"]
        self.assertEqual(asn["outcome"], "assigned")
        self.assertIsNotNone(asn["assignment_id"])
        self.assertIsNotNone(asn["lease_id"])
        self.assertEqual(asn["work_number"], 600)
        self.assertIn("implement", asn["allowed_actions"])
        self.assertIn("merge", asn["forbidden_actions"])
        proof = res["lease_proof"]
        self.assertEqual(proof["source"], "control_plane_db.assign_and_lease")

    def test_concurrent_allocators_no_double_assign(self) -> None:
        cand = WorkCandidate(
            kind="pr",
            number=100,
            head_sha="a" * 40,
            priority=10,
        )
        # Seed work item so both race the same target
        self.db.upsert_work_item(
            remote="prgs",
            org="org",
            repo="repo",
            kind="pr",
            number=100,
            current_head_sha="a" * 40,
        )
        results = []
        lock = threading.Lock()

        def worker(sid: str):
            r = allocate_next_work(
                self.db,
                session_id=sid,
                role="reviewer",
                remote="prgs",
                org="org",
                repo="repo",
                candidates=[cand],
                apply=True,
                profile_name="prgs-reviewer",
            )
            with lock:
                results.append(r)

        with ThreadPoolExecutor(max_workers=2) as pool:
            futs = [pool.submit(worker, f"s-{i}") for i in range(2)]
            for f in as_completed(futs):
                f.result()

        outcomes = [r["outcome"] for r in results]
        self.assertEqual(outcomes.count(OUTCOME_ASSIGNED), 1, outcomes)
        self.assertEqual(outcomes.count(OUTCOME_WAIT), 1, outcomes)

    def test_blocked_and_dependency_skipped(self) -> None:
        cands = [
            WorkCandidate(kind="issue", number=1, blocked=True, labels=("status:blocked",)),
            WorkCandidate(
                kind="issue",
                number=612,
                labels=("status:ready",),
                dependency_unmet=True,
                dependency_reason="downstream of #600",
            ),
        ]
        res = self._alloc(candidates=cands, apply=True, role="author")
        self.assertEqual(res["outcome"], OUTCOME_NO_SAFE)
        self.assertIsNone(res["selected"])
        reasons = " ".join(s["reason"] for s in res["skipped"])
        self.assertIn("blocked", reasons.lower())
        self.assertIn("600", reasons)

    def test_already_leased_returns_wait(self) -> None:
        self.db.upsert_session(session_id="owner", role="author")
        self.db.assign_and_lease(
            session_id="owner",
            role="author",
            remote="prgs",
            org="org",
            repo="repo",
            kind="issue",
            number=50,
        )
        res = self._alloc(
            session_id="other",
            candidates=[
                WorkCandidate(kind="issue", number=50, labels=("status:ready",))
            ],
            apply=True,
        )
        self.assertEqual(res["outcome"], OUTCOME_WAIT)
        self.assertEqual(res["owner_session_id"], "owner")

    def test_role_ineligible_skipped(self) -> None:
        # PR with current-head REQUEST_CHANGES expects author, not reviewer.
        cand = WorkCandidate(
            kind="pr",
            number=9,
            head_sha="b" * 40,
            request_changes_current_head=True,
        )
        self.assertEqual(expected_role_for_candidate(cand), "author")
        res = self._alloc(
            role="reviewer",
            profile_name="prgs-reviewer",
            candidates=[cand],
            apply=False,
        )
        self.assertEqual(res["outcome"], OUTCOME_NO_SAFE)
        self.assertTrue(any("expects role" in s["reason"] for s in res["skipped"]))

    def test_terminal_lock_blocks_downstream_reviewer_prs(self) -> None:
        self.db.set_terminal_lock(
            remote="prgs",
            org="org",
            repo="repo",
            terminal_pr=10,
            decision="request_changes",
            status="active",
        )
        cands = [
            WorkCandidate(kind="pr", number=11, head_sha="c" * 40, priority=5),
            WorkCandidate(kind="pr", number=10, head_sha="d" * 40, priority=1),
        ]
        res = self._alloc(
            role="reviewer",
            profile_name="prgs-reviewer",
            candidates=cands,
            apply=False,
        )
        # Terminal PR #10 is selectable; #11 skipped for terminal path.
        self.assertEqual(res["selected"]["number"], 10)
        self.assertTrue(
            any(
                s["number"] == 11 and "terminal-review lock" in s["reason"]
                for s in res["skipped"]
            )
        )

    def test_terminal_lock_blocks_all_when_only_downstream(self) -> None:
        self.db.set_terminal_lock(
            remote="prgs",
            org="org",
            repo="repo",
            terminal_pr=10,
            decision="request_changes",
            status="active",
        )
        res = self._alloc(
            role="reviewer",
            profile_name="prgs-reviewer",
            candidates=[
                WorkCandidate(kind="pr", number=99, head_sha="e" * 40),
            ],
            apply=False,
        )
        self.assertEqual(res["outcome"], OUTCOME_BLOCKED_TERMINAL)

    def test_no_work_structured(self) -> None:
        res = self._alloc(candidates=[], apply=True)
        self.assertEqual(res["outcome"], OUTCOME_NO_SAFE)
        self.assertIsNone(res["selected"])
        self.assertTrue(res["reasons"])

    def test_rejects_incident_kind(self) -> None:
        with self.assertRaises(InvalidWorkKindError):
            WorkCandidate(kind="sentry_incident", number=1)

    def test_612_downstream_marker_in_result(self) -> None:
        res = self._alloc(
            candidates=[
                WorkCandidate(kind="issue", number=600, labels=("status:ready",))
            ],
            apply=True,
        )
        self.assertIn("612", res.get("downstream_note", ""))

    def test_substrate_not_file_or_comment_lease(self) -> None:
        res = self._alloc(
            candidates=[
                WorkCandidate(kind="issue", number=1, labels=("status:ready",))
            ],
            apply=True,
        )
        self.assertEqual(res["substrate"], "control_plane_db")
        self.assertFalse(res["file_lock_only"])
        self.assertFalse(res["comment_lease_only"])
        self.assertEqual(
            res["lease_proof"]["source"], "control_plane_db.assign_and_lease"
        )

    def test_candidate_from_dict(self) -> None:
        c = candidate_from_dict(
            {
                "kind": "pr",
                "number": 7,
                "head_sha": "f" * 40,
                "approval_on_current_head": True,
                "mergeable": True,
            }
        )
        self.assertEqual(expected_role_for_candidate(c), "merger")

    def test_merger_gets_clean_approval(self) -> None:
        c = WorkCandidate(
            kind="pr",
            number=3,
            head_sha="1" * 40,
            approval_on_current_head=True,
            mergeable=True,
            priority=100,
        )
        res = self._alloc(
            role="merger",
            profile_name="prgs-merger",
            candidates=[c],
            apply=True,
            session_id="merger-1",
        )
        self.assertEqual(res["outcome"], OUTCOME_ASSIGNED)
        self.assertEqual(res["assignment"]["work_number"], 3)
        self.assertIn("merge", res["assignment"]["allowed_actions"])

    def test_db_unavailable_fails_closed(self) -> None:
        res = allocate_next_work(
            None,  # type: ignore[arg-type]
            session_id="x",
            role="author",
            remote="prgs",
            org="o",
            repo="r",
            candidates=[],
            apply=True,
        )
        self.assertFalse(res["success"])
        self.assertIn("unavailable", res["reasons"][0].lower())


class SideEffectFreeAllocationTest(unittest.TestCase):
    """``side_effect_free`` dry runs write nothing to the control plane (#643).

    A plain ``apply=False`` still called ``upsert_session`` and
    ``expire_stale_leases`` before the apply branch was consulted, so a caller
    advertising a read-only preview mutated on every call — one unreferenced
    session row per preview, plus a global lease sweep.
    """

    def setUp(self) -> None:
        self._tmp = tempfile.TemporaryDirectory()
        self.db = ControlPlaneDB(os.path.join(self._tmp.name, "cp.sqlite3"))

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

    def _alloc(self, **kwargs):
        defaults = dict(
            db=self.db,
            session_id="s-preview",
            role="author",
            remote="prgs",
            org="org",
            repo="repo",
            candidates=[
                WorkCandidate(kind="issue", number=643, labels=("status:ready",))
            ],
            apply=False,
            profile_name="prgs-author",
            username="jcwalker3",
        )
        defaults.update(kwargs)
        return allocate_next_work(**defaults)

    def _session_ids(self) -> set[str]:
        return {str(r.get("session_id")) for r in self.db.list_sessions()}

    def test_side_effect_free_preview_writes_no_session_row(self):
        before = self._session_ids()
        result = self._alloc(side_effect_free=True)
        self.assertEqual(result["outcome"], OUTCOME_PREVIEW)
        self.assertEqual(self._session_ids(), before)
        self.assertNotIn("s-preview", self._session_ids())

    def test_plain_dry_run_still_registers_a_session(self):
        # The default is unchanged for every existing caller.
        self._alloc()
        self.assertIn("s-preview", self._session_ids())

    def test_repeated_previews_do_not_accumulate_rows(self):
        for index in range(5):
            self._alloc(side_effect_free=True, session_id=f"s-{index}")
        self.assertEqual(self._session_ids(), set())

    def test_side_effect_free_does_not_sweep_stale_leases(self):
        self.db.upsert_session(session_id="owner", role="author", pid=1)
        assigned = self.db.assign_and_lease(
            session_id="owner",
            role="author",
            remote="prgs",
            org="org",
            repo="repo",
            kind="issue",
            number=999,
            lease_ttl_seconds=-60,  # already expired
        )
        self.assertEqual(assigned.outcome, "assigned")

        self._alloc(side_effect_free=True)

        # The expired row is still 'active' in the DB: nothing swept it.
        statuses = {
            r["lease_id"]: r["status"]
            for r in self.db.list_leases(
                remote="prgs", org="org", repo="repo",
                statuses=("active", "expired"),
            )
        }
        self.assertEqual(statuses.get(assigned.lease_id), "active")

    def test_expired_claims_are_filtered_in_memory_so_work_stays_selectable(self):
        """The read-only mirror of the sweep: expired claims must not block."""
        self.db.upsert_session(session_id="owner", role="author", pid=1)
        self.db.assign_and_lease(
            session_id="owner",
            role="author",
            remote="prgs",
            org="org",
            repo="repo",
            kind="issue",
            number=643,
            lease_ttl_seconds=-60,  # expired: must not withhold #643
        )
        result = self._alloc(side_effect_free=True)
        self.assertEqual(result["outcome"], OUTCOME_PREVIEW)
        self.assertEqual(result["selected"]["number"], 643)

    def test_a_live_claim_still_withholds_the_work(self):
        self.db.upsert_session(session_id="owner", role="author", pid=1)
        self.db.assign_and_lease(
            session_id="owner",
            role="author",
            remote="prgs",
            org="org",
            repo="repo",
            kind="issue",
            number=643,
            lease_ttl_seconds=3600,
        )
        result = self._alloc(side_effect_free=True)
        self.assertNotEqual(result["outcome"], OUTCOME_ASSIGNED)
        self.assertNotEqual((result.get("selected") or {}).get("number"), 643)

    def test_side_effect_free_with_apply_fails_closed(self):
        result = self._alloc(side_effect_free=True, apply=True)
        self.assertFalse(result["success"])
        self.assertEqual(result["outcome"], OUTCOME_NO_SAFE)
        self.assertIsNone(result["assignment"])
        self.assertIn("incompatible with apply", result["reasons"][0])
        # And it reserved nothing.
        self.assertEqual(
            self.db.list_leases(remote="prgs", org="org", repo="repo"), []
        )


class DropExpiredClaimsTest(unittest.TestCase):
    """The in-memory expiry filter behind side-effect-free previews (#643)."""

    def test_unparseable_expiry_is_kept_rather_than_assumed_free(self):
        claims = {
            ("issue", 1): {"lease_id": "l1", "expires_at": "not-a-date"},
            ("issue", 2): {"lease_id": "l2"},
            ("issue", 3): {"lease_id": "l3", "expires_at": None},
        }
        self.assertEqual(_drop_expired_claims(claims), claims)

    def test_expired_dropped_and_future_kept(self):
        now = datetime(2026, 7, 25, 12, 0, tzinfo=timezone.utc)
        claims = {
            ("issue", 1): {"expires_at": "2026-07-25T11:59:59+00:00"},
            ("issue", 2): {"expires_at": "2026-07-25T12:00:01+00:00"},
            ("issue", 3): {"expires_at": "2026-07-25T12:00:00+00:00"},  # boundary
        }
        kept = _drop_expired_claims(claims, now=now)
        self.assertEqual(set(kept), {("issue", 2)})

    def test_naive_and_zulu_timestamps_are_treated_as_utc(self):
        now = datetime(2026, 7, 25, 12, 0, tzinfo=timezone.utc)
        claims = {
            ("issue", 1): {"expires_at": "2026-07-25T11:00:00"},  # naive, past
            ("issue", 2): {"expires_at": "2026-07-25T13:00:00Z"},  # zulu, future
        }
        kept = _drop_expired_claims(claims, now=now)
        self.assertEqual(set(kept), {("issue", 2)})


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