"""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 allocator_service import (
    OUTCOME_ASSIGNED,
    OUTCOME_BLOCKED_TERMINAL,
    OUTCOME_NO_SAFE,
    OUTCOME_PREVIEW,
    OUTCOME_WAIT,
    WorkCandidate,
    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())


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