"""Ownership guards at scheduler mutation/recovery/dispatch entries (R4/F5).

The held lock must be proven at the entry points, not merely exist as a
class: ``poll``, schedule administration, and ``recover`` reject when the
caller cannot prove live ownership, before any write or side effect.
"""

from __future__ import annotations

from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any

import pytest

from tests.scheduling.controlled import (
    DictDeployments,
    ScriptedDispatcher,
    fixture_environment,
)
from wf_artifacts.runs.store import FileRunStore
from wf_scheduling import recovery as sched_recovery
from wf_scheduling.calendar import OneShotSource
from wf_scheduling.models import Schedule
from wf_scheduling.ownership import SchedulerOwnership, SecondOwnerError
from wf_scheduling.poll import Scheduler
from wf_scheduling.prepare import SchedulePreparer
from wf_scheduling.store import FileScheduleStore


def ts(y: int, mo: int, d: int, h: int = 0, mi: int = 0) -> datetime:
    return datetime(y, mo, d, h, mi, tzinfo=UTC)


def _sched_model(sid: str, **kw: Any) -> Schedule:
    now = ts(2026, 9, 8, 12, 0)
    base: dict[str, Any] = {
        "id": sid,
        "deployment_id": "dep-1",
        "trigger": {"kind": "cron", "expression": "0 * * * *", "timezone": "UTC"},
        "input_bindings": [],
        "created_at": now.isoformat(),
        "updated_at": now.isoformat(),
    }
    base.update(kw)
    return Schedule.model_validate(base)


def _scheduler(
    tmp_path: Path,
    ownership: SchedulerOwnership,
    *,
    script: dict | None = None,
) -> tuple[Scheduler, FileScheduleStore, FileRunStore]:
    sched_store = FileScheduleStore(tmp_path / "sched")
    run_store = FileRunStore(tmp_path / "runs")
    sched = Scheduler(
        schedule_store=sched_store,
        run_store=run_store,
        sources={},
        capacity=4,
        preparer=SchedulePreparer(
            DictDeployments({"dep-1": {"rev": 1, "required": []}}),
            fixture_environment,
        ),
        dispatcher=ScriptedDispatcher(script),
        ownership=ownership,
    )
    return sched, sched_store, run_store


def _due_setup(sched: Scheduler, store: FileScheduleStore, intended: datetime) -> None:
    store.create_schedule(_sched_model("a"))
    store.save_consumed("a", intended - timedelta(hours=1))
    sched.sources["a"] = OneShotSource(intended)


def test_poll_without_ownership_rejects_before_writes(tmp_path: Path) -> None:
    calls: list[str] = []

    def spy(admission: Any, now: datetime) -> Any:
        calls.append(admission.id)
        raise AssertionError("dispatcher must not run without ownership")

    ownership = SchedulerOwnership(tmp_path / "sched", owner="never-acquired")
    sched, store, runs = _scheduler(tmp_path, ownership)
    sched.dispatcher = ScriptedDispatcher({"*": spy})
    intended = ts(2026, 9, 8, 12, 0)
    _due_setup(sched, store, intended)
    with pytest.raises(SecondOwnerError):
        sched.poll(intended)
    assert calls == []
    assert runs.list_runs() == []
    assert runs.list_admissions() == []
    assert store.list_occurrences("a", limit=100)["total"] == 0
    assert store.get_consumed("a") == intended - timedelta(hours=1)


def test_poll_with_released_ownership_rejects(tmp_path: Path) -> None:
    ownership = SchedulerOwnership(tmp_path / "sched", owner="test").acquire()
    sched, store, runs = _scheduler(tmp_path, ownership)
    intended = ts(2026, 9, 8, 12, 0)
    _due_setup(sched, store, intended)
    ownership.release()
    assert not ownership.held
    with pytest.raises(SecondOwnerError):
        sched.poll(intended)
    assert runs.list_runs() == []
    assert store.list_occurrences("a", limit=100)["total"] == 0


def test_recover_without_ownership_rejects_before_writes(tmp_path: Path) -> None:
    from tests.artifacts.test_run_store import artifact as _artifact
    from tests.artifacts.test_run_store import deployment as _deployment
    from wf_api.run_lifecycle import materialize_admitted_view, persist_admission
    from wf_artifacts import PinnedRunEnvironment

    sched_store = FileScheduleStore(tmp_path / "sched")
    run_store = FileRunStore(tmp_path / "runs")
    sched_store.create_schedule(_sched_model("a"))
    env = PinnedRunEnvironment(
        deployment=_deployment(), root_artifact=_artifact(), child_artifacts=[]
    )
    admission = persist_admission(
        store=run_store,
        run_id=run_store.allocate_run_id(),
        environment=env,
        resolved_input={},
        max_steps=None,
        scheduled_at=ts(2026, 9, 8, 12, 0),
        schedule_id="a",
        schedule_revision=1,
    )
    materialize_admitted_view(store=run_store, admission=admission)
    ownership = SchedulerOwnership(tmp_path / "sched", owner="never-acquired")
    with pytest.raises(SecondOwnerError):
        sched_recovery.recover(
            schedule_store=sched_store,
            run_store=run_store,
            now=ts(2026, 9, 8, 12, 0),
            ownership=ownership,
        )
    # Rejected before any write: the admitted run is untouched.
    assert run_store.get_run(admission.id).status.value == "admitted"


def test_held_ownership_allows_poll_and_recovery(tmp_path: Path) -> None:
    ownership = SchedulerOwnership(tmp_path / "sched", owner="test").acquire()
    try:
        sched, store, runs = _scheduler(tmp_path, ownership, script={"*": "hang"})
        intended = ts(2026, 9, 8, 12, 0)
        _due_setup(sched, store, intended)
        assert sched.poll(intended) == {"a": f"admit:{runs.list_runs()[0].id}"}
        diags = sched_recovery.recover(
            schedule_store=store,
            run_store=runs,
            now=intended,
            ownership=ownership,
        )
        assert diags != []
    finally:
        ownership.release()


def test_second_owner_cannot_acquire_for_poll(tmp_path: Path) -> None:
    first = SchedulerOwnership(tmp_path / "sched", owner="first").acquire()
    try:
        second = SchedulerOwnership(tmp_path / "sched", owner="second")
        with pytest.raises(SecondOwnerError):
            second.acquire()
        assert not second.held
        # A scheduler proving the loser's (unheld) ownership is rejected.
        sched, store, _ = _scheduler(tmp_path, second)
        intended = ts(2026, 9, 8, 12, 0)
        _due_setup(sched, store, intended)
        with pytest.raises(SecondOwnerError):
            sched.poll(intended)
    finally:
        first.release()


def test_admin_mutations_require_ownership(tmp_path: Path) -> None:
    ownership = SchedulerOwnership(tmp_path / "sched", owner="test")
    sched, store, _ = _scheduler(tmp_path, ownership)
    store.create_schedule(_sched_model("a", paused=True))
    with pytest.raises(SecondOwnerError):
        sched.resume_schedule("a", ts(2026, 9, 8, 12, 0))
    with pytest.raises(SecondOwnerError):
        sched.edit_schedule("a", ts(2026, 9, 8, 12, 0))
    assert store.get_schedule("a").paused is True
    assert store.get_schedule("a").revision == 1