"""Allocator ownership exclusion tests (#765).

One session's active lease must never blockade the author queue for a
different controller. Covers: foreign lease skipped, next unclaimed candidate
selected, own task resumable, task-local blocker quarantined, all-claimed ->
wait, same profile + different controller_instance_id -> different ownership,
and claimed candidates reported in skipped results.
"""

from __future__ import annotations

import os
import tempfile
import unittest

from allocator_service import (
    OUTCOME_OWNERSHIP_DEFECT,
    OUTCOME_PREVIEW,
    OUTCOME_WAIT,
    OWNERSHIP_FOREIGN,
    OWNERSHIP_OWN,
    OWNERSHIP_UNKNOWN,
    SKIP_CLAIMED_BY_OTHER_SESSION,
    WorkCandidate,
    allocate_next_work,
    classify_claim_ownership,
    resolve_controller_instance_id,
)
from control_plane_db import ControlPlaneDB

REMOTE = "prgs"
ORG = "Scaled-Tech-Consulting"
REPO = "Gitea-Tools"

MINE = "ctl-mine-0001"
THEIRS = "ctl-theirs-0002"


def _issue(number: int, **kwargs) -> WorkCandidate:
    base = dict(
        kind="issue",
        number=number,
        state="open",
        labels=("status:ready", "type:bug"),
        title=f"issue {number}",
        priority=20,
    )
    base.update(kwargs)
    return WorkCandidate(**base)


def _claim(number: int, *, session_id: str, instance: str | None, kind: str = "issue"):
    return {
        "lease_id": f"lease-{number}",
        "session_id": session_id,
        "controller_instance_id": instance,
        "role": "author",
        "profile": "prgs-author",
        "expires_at": "2026-07-20T07:06:09Z",
        "work_kind": kind,
        "work_number": number,
    }


class AllocatorOwnershipTestCase(unittest.TestCase):
    def setUp(self):
        self._tmp = tempfile.TemporaryDirectory()
        self.addCleanup(self._tmp.cleanup)
        self.db = ControlPlaneDB(os.path.join(self._tmp.name, "cp.sqlite3"))

    def _allocate(
        self,
        candidates,
        *,
        claims,
        session_id="sess-mine",
        instance=MINE,
        apply=False,
        role="author",
    ):
        return allocate_next_work(
            self.db,
            session_id=session_id,
            role=role,
            remote=REMOTE,
            org=ORG,
            repo=REPO,
            candidates=candidates,
            apply=apply,
            profile_name="prgs-author",
            controller_instance_id=instance,
            claims=claims,
        )


class TestOwnershipClassification(AllocatorOwnershipTestCase):
    def test_no_claim_returns_none(self):
        self.assertIsNone(
            classify_claim_ownership(
                None, session_id="s", controller_instance_id=MINE
            )
        )

    def test_same_controller_instance_is_own(self):
        claim = _claim(1, session_id="other-session", instance=MINE)
        self.assertEqual(
            classify_claim_ownership(
                claim, session_id="sess-mine", controller_instance_id=MINE
            ),
            OWNERSHIP_OWN,
        )

    def test_same_profile_different_instance_is_foreign(self):
        """Shared profile must not imply shared ownership."""
        claim = _claim(1, session_id="other-session", instance=THEIRS)
        self.assertEqual(
            classify_claim_ownership(
                claim, session_id="sess-mine", controller_instance_id=MINE
            ),
            OWNERSHIP_FOREIGN,
        )

    def test_exact_session_match_is_own(self):
        claim = _claim(1, session_id="sess-mine", instance=None)
        self.assertEqual(
            classify_claim_ownership(
                claim, session_id="sess-mine", controller_instance_id=None
            ),
            OWNERSHIP_OWN,
        )

    def test_legacy_claim_with_neither_side_identified_is_foreign(self):
        """No identities anywhere: a different session id is simply not ours."""
        claim = _claim(1, session_id="someone-else", instance=None)
        self.assertEqual(
            classify_claim_ownership(
                claim, session_id="sess-mine", controller_instance_id=None
            ),
            OWNERSHIP_FOREIGN,
        )

    def test_claim_identified_but_local_undeclared_is_unknown(self):
        """Only one side identified: not comparable, so never adopt."""
        claim = _claim(1, session_id="someone-else", instance=THEIRS)
        self.assertEqual(
            classify_claim_ownership(
                claim, session_id="sess-mine", controller_instance_id=None
            ),
            OWNERSHIP_UNKNOWN,
        )

    def test_local_identified_but_claim_undeclared_is_unknown(self):
        """A legacy lease may be our own under an old session id; do not guess."""
        claim = _claim(1, session_id="someone-else", instance=None)
        self.assertEqual(
            classify_claim_ownership(
                claim, session_id="sess-mine", controller_instance_id=MINE
            ),
            OWNERSHIP_UNKNOWN,
        )

    def test_resolve_controller_instance_id_reads_env(self):
        self.assertEqual(
            resolve_controller_instance_id({"GITEA_CONTROLLER_INSTANCE_ID": MINE}),
            MINE,
        )
        self.assertIsNone(resolve_controller_instance_id({}))
        self.assertIsNone(
            resolve_controller_instance_id({"GITEA_CONTROLLER_INSTANCE_ID": "  "})
        )


class TestForeignLeaseDoesNotBlockade(AllocatorOwnershipTestCase):
    def test_foreign_claim_skipped_and_next_issue_selected(self):
        """Skip the claimed issue, select the next unclaimed one."""
        candidates = [_issue(607), _issue(615), _issue(617)]
        claims = {
            ("issue", 607): _claim(607, session_id="sess-theirs", instance=THEIRS)
        }
        result = self._allocate(candidates, claims=claims)

        self.assertEqual(result["outcome"], OUTCOME_PREVIEW)
        self.assertEqual(result["selected"]["number"], 615)
        skipped_607 = [s for s in result["skipped"] if s["number"] == 607]
        self.assertEqual(len(skipped_607), 1)
        self.assertEqual(
            skipped_607[0]["reason_code"], SKIP_CLAIMED_BY_OTHER_SESSION
        )
        self.assertIn(SKIP_CLAIMED_BY_OTHER_SESSION, skipped_607[0]["reason"])

    def test_claimed_candidate_appears_in_skipped_inventory(self):
        """Skipped reporting must reflect claimed candidates."""
        candidates = [_issue(607), _issue(615)]
        claims = {
            ("issue", 607): _claim(607, session_id="sess-theirs", instance=THEIRS)
        }
        result = self._allocate(candidates, claims=claims)
        self.assertEqual(len(result["skipped"]), 1)
        self.assertEqual(len(result["claims_excluded"]), 1)
        excluded = result["claims_excluded"][0]
        self.assertEqual(excluded["number"], 607)
        self.assertEqual(excluded["ownership"], OWNERSHIP_FOREIGN)
        self.assertEqual(excluded["owner_controller_instance_id"], THEIRS)

    def test_task_local_blocker_does_not_freeze_unrelated_work(self):
        """A quarantined task must not stop the rest of the queue."""
        candidates = [_issue(607), _issue(615), _issue(617)]
        claims = {
            ("issue", 607): _claim(607, session_id="sess-theirs", instance=THEIRS)
        }
        first = self._allocate(candidates, claims=claims)
        self.assertEqual(first["selected"]["number"], 615)

        # 615 then gets claimed by yet another controller; queue still advances.
        claims[("issue", 615)] = _claim(
            615, session_id="sess-third", instance="ctl-third-0003"
        )
        second = self._allocate(candidates, claims=claims)
        self.assertEqual(second["selected"]["number"], 617)

    def test_multiple_controllers_get_different_issues(self):
        """Concurrent author sessions work on different issues."""
        candidates = [_issue(607), _issue(615)]
        claims = {
            ("issue", 607): _claim(607, session_id="sess-theirs", instance=THEIRS)
        }
        mine = self._allocate(candidates, claims=claims, instance=MINE)
        theirs = self._allocate(
            candidates, claims=claims, session_id="sess-theirs", instance=THEIRS
        )
        self.assertEqual(mine["selected"]["number"], 615)
        # The other controller may still be handed its own in-progress task.
        self.assertEqual(theirs["selected"]["number"], 607)

    def test_unclaimed_queue_is_unaffected(self):
        candidates = [_issue(607), _issue(615)]
        result = self._allocate(candidates, claims={})
        self.assertEqual(result["selected"]["number"], 607)
        self.assertEqual(result["skipped"], [])
        self.assertEqual(result["claims_excluded"], [])


class TestOwnTaskResume(AllocatorOwnershipTestCase):
    def test_controller_may_resume_its_own_active_task(self):
        """Own claim stays selectable across a new session id."""
        candidates = [_issue(607), _issue(615)]
        claims = {
            ("issue", 607): _claim(607, session_id="sess-mine-old", instance=MINE)
        }
        result = self._allocate(
            candidates, claims=claims, session_id="sess-mine-new", instance=MINE
        )
        self.assertEqual(result["selected"]["number"], 607)
        self.assertEqual(result["claims_excluded"], [])

    def test_own_claim_by_exact_session_is_selectable(self):
        candidates = [_issue(607)]
        claims = {("issue", 607): _claim(607, session_id="sess-mine", instance=None)}
        result = self._allocate(
            candidates, claims=claims, session_id="sess-mine", instance=None
        )
        self.assertEqual(result["selected"]["number"], 607)


class TestAllCandidatesClaimed(AllocatorOwnershipTestCase):
    def test_all_claimed_returns_wait_not_a_claimed_selection(self):
        """Never hand back a claimed issue; report waiting instead."""
        candidates = [_issue(607), _issue(615)]
        claims = {
            ("issue", 607): _claim(607, session_id="sess-a", instance=THEIRS),
            ("issue", 615): _claim(615, session_id="sess-b", instance="ctl-c-0003"),
        }
        result = self._allocate(candidates, claims=claims)
        self.assertIsNone(result["selected"])
        self.assertEqual(result["outcome"], OUTCOME_WAIT)
        self.assertEqual(len(result["claims_excluded"]), 2)

    def test_unidentifiable_owner_reports_ownership_defect(self):
        """Refuse to adopt when ownership cannot be established."""
        candidates = [_issue(607)]
        claims = {("issue", 607): _claim(607, session_id="sess-legacy", instance=None)}
        result = self._allocate(candidates, claims=claims)
        self.assertIsNone(result["selected"])
        self.assertEqual(result["outcome"], OUTCOME_OWNERSHIP_DEFECT)
        self.assertEqual(len(result["ownership_defects"]), 1)
        self.assertEqual(
            result["ownership_defects"][0]["ownership"], OWNERSHIP_UNKNOWN
        )


class TestClaimsFromControlPlaneDb(AllocatorOwnershipTestCase):
    """End-to-end against the real substrate, not injected claim dicts."""

    def _seed_lease(self, number: int, *, session_id: str, instance: str | None):
        self.db.upsert_session(
            session_id=session_id,
            role="author",
            profile="prgs-author",
            pid=4242,
            controller_instance_id=instance,
        )
        return self.db.assign_and_lease(
            session_id=session_id,
            role="author",
            remote=REMOTE,
            org=ORG,
            repo=REPO,
            kind="issue",
            number=number,
        )

    def test_controller_instance_id_persists_on_session(self):
        row = self.db.upsert_session(
            session_id="sess-x",
            role="author",
            profile="prgs-author",
            pid=1,
            controller_instance_id=MINE,
        )
        self.assertEqual(row["controller_instance_id"], MINE)

    def test_heartbeat_without_instance_does_not_erase_ownership(self):
        self.db.upsert_session(
            session_id="sess-x",
            role="author",
            profile="prgs-author",
            pid=1,
            controller_instance_id=MINE,
        )
        row = self.db.upsert_session(
            session_id="sess-x", role="author", profile="prgs-author", pid=1
        )
        self.assertEqual(row["controller_instance_id"], MINE)

    def test_list_active_claims_surfaces_owner_instance(self):
        self._seed_lease(607, session_id="sess-theirs", instance=THEIRS)
        claims = self.db.list_active_claims(remote=REMOTE, org=ORG, repo=REPO)
        self.assertIn(("issue", 607), claims)
        self.assertEqual(claims[("issue", 607)]["controller_instance_id"], THEIRS)

    def test_live_foreign_lease_is_excluded_without_injected_claims(self):
        self._seed_lease(607, session_id="sess-theirs", instance=THEIRS)
        result = allocate_next_work(
            self.db,
            session_id="sess-mine",
            role="author",
            remote=REMOTE,
            org=ORG,
            repo=REPO,
            candidates=[_issue(607), _issue(615)],
            apply=False,
            profile_name="prgs-author",
            controller_instance_id=MINE,
        )
        self.assertEqual(result["selected"]["number"], 615)
        self.assertEqual(
            result["skipped"][0]["reason_code"], SKIP_CLAIMED_BY_OTHER_SESSION
        )

    def test_apply_reserves_the_unclaimed_issue(self):
        self._seed_lease(607, session_id="sess-theirs", instance=THEIRS)
        result = allocate_next_work(
            self.db,
            session_id="sess-mine",
            role="author",
            remote=REMOTE,
            org=ORG,
            repo=REPO,
            candidates=[_issue(607), _issue(615)],
            apply=True,
            profile_name="prgs-author",
            controller_instance_id=MINE,
        )
        self.assertEqual(result["outcome"], "assigned_work")
        self.assertEqual(result["selected"]["number"], 615)
        self.assertEqual(result["assignment"]["work_number"], 615)


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