"""Durable dispatch transition: executing is persisted BEFORE the executor (R4/F1).

Pending means provably undispatched and safe to dispatch. The scheduler
persists the undispatched→executing transition before invoking the
executor; a crash afterwards abandons the run (failed, no replay) instead
of retrying. ``finally`` cannot be trusted across process termination, so
marker clearing is explicit on the stopped path only.
"""

from __future__ import annotations

import multiprocessing as mp
from datetime import UTC, datetime, timedelta
from pathlib import Path
from typing import Any, cast

import pytest

from tests.scheduling.controlled import (
    DictDeployments,
    ScriptedDispatcher,
    fixture_environment,
    stopped_state,
)
from tests.scheduling.controlled import ScriptedDispatcher as SD
from wf_artifacts.runs.store import FileRunStore
from wf_scheduling import recovery as sched_recovery
from wf_scheduling.calendar import OneShotSource
from wf_scheduling.dispatch import StillRunning, Stopped
from wf_scheduling.models import Schedule
from wf_scheduling.ownership import SchedulerOwnership
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 _harness(
    root: Path,
    ownership: SchedulerOwnership,
    *,
    capacity: int = 4,
    script: dict | None = None,
) -> tuple[Scheduler, FileScheduleStore, FileRunStore]:
    sched_store = FileScheduleStore(root / "sched")
    run_store = FileRunStore(root / "runs")
    sched = Scheduler(
        schedule_store=sched_store,
        run_store=run_store,
        sources={},
        capacity=capacity,
        preparer=SchedulePreparer(
            DictDeployments({"dep-1": {"rev": 1, "required": []}}),
            fixture_environment,
        ),
        dispatcher=ScriptedDispatcher(script),
        ownership=ownership,
    )
    return sched, sched_store, run_store


def _due(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_executing_is_set_before_executor_runs(tmp_path: Path) -> None:
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        observed: dict[str, bool] = {}

        def probe(admission: Any, now: datetime) -> Any:
            observed["executing_at_entry"] = FileRunStore(
                tmp_path / "runs"
            ).is_executing(admission.id)
            observed["pending_at_entry"] = FileRunStore(
                tmp_path / "runs"
            ).is_pending_dispatch(admission.id)
            return StillRunning()

        sched, store, runs = _harness(tmp_path, ownership, script={"*": probe})
        intended = ts(2026, 9, 8, 12, 0)
        _due(sched, store, intended)
        sched.poll(intended)
        assert observed == {"executing_at_entry": True, "pending_at_entry": False}
        run_id = runs.list_runs()[0].id
        assert runs.is_executing(run_id)
        assert not runs.is_pending_dispatch(run_id)
    finally:
        ownership.release()


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

    class FailMarkOnce(FileRunStore):
        def __init__(self, root: Path) -> None:
            super().__init__(root)
            self.armed = True

        def mark_executing(self, run_id: str) -> None:
            if self.armed:
                self.armed = False
                raise OSError("injected mark_executing failure")
            super().mark_executing(run_id)

    def spy(admission: Any, now: datetime) -> Any:
        calls.append(admission.id)
        return StillRunning()

    sched_store = FileScheduleStore(tmp_path / "sched")
    run_store = FailMarkOnce(tmp_path / "runs")
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        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({"*": spy}),
            ownership=ownership,
        )
        intended = ts(2026, 9, 8, 12, 0)
        _due(sched, sched_store, intended)
        with pytest.raises(OSError, match="injected mark_executing"):
            sched.poll(intended)
        # The executor never ran; the run stays provably undispatched.
        assert calls == []
        run_id = run_store.list_runs()[0].id
        assert run_store.is_pending_dispatch(run_id)
        assert not run_store.is_executing(run_id)
        # Next poll dispatches exactly once.
        sched.poll(intended + timedelta(minutes=1))
        assert calls == [run_id]
    finally:
        ownership.release()


def test_capacity_shortage_keeps_pending_undispatched(tmp_path: Path) -> None:
    calls: list[str] = []
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        sched, store, runs = _harness(tmp_path, ownership, capacity=0)
        intended = ts(2026, 9, 8, 12, 0)
        _due(sched, store, intended)

        def spy(admission: Any, now: datetime) -> Any:
            calls.append(admission.id)
            return StillRunning()

        cast(SD, sched.dispatcher).script = {"*": spy}
        assert sched.poll(intended) == {"a": "admit:held-undecided"}
        assert calls == []
        # A held candidate is pending-but-unadmitted, not an admitted run.
        assert runs.list_runs() == []
        sched.capacity = 4
        result = sched.poll(intended + timedelta(minutes=1))
        assert result["a"].startswith("admit:run-")
        assert len(calls) == 1
    finally:
        ownership.release()


def test_pending_on_terminal_run_clears_without_redispatch(tmp_path: Path) -> None:
    calls: list[str] = []
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        sched, store, runs = _harness(tmp_path, ownership)

        def spy(admission: Any, now: datetime) -> Any:
            calls.append(admission.id)
            return Stopped(result=stopped_state(admission, "complete"))

        cast(SD, sched.dispatcher).script = {"*": spy}
        intended = ts(2026, 9, 8, 12, 0)
        _due(sched, store, intended)
        sched.poll(intended)
        assert calls != []
        run_id = runs.list_runs()[0].id
        assert runs.get_run(run_id).status.value == "completed"
        # A stale pending marker on a terminal run must not redispatch.
        runs.mark_pending_dispatch(run_id)
        before = len(calls)
        sched.poll(intended + timedelta(minutes=5))
        assert len(calls) == before
        assert not runs.is_pending_dispatch(run_id)
        assert runs.get_run(run_id).status.value == "completed"
        page = store.list_occurrences("a", limit=100)
        completed = [
            r
            for r in cast(list[dict[str, Any]], page["occurrences"])
            if r["kind"] == "completed"
        ]
        assert len(completed) == 1
    finally:
        ownership.release()


def test_pending_without_admission_fails_closed(tmp_path: Path) -> None:
    from datetime import datetime as _dt

    from tests.artifacts.test_run_store import artifact as _artifact
    from tests.artifacts.test_run_store import deployment as _deployment
    from wf_artifacts import PinnedRunEnvironment
    from wf_artifacts.runs.models import (
        ResumeReadiness,
        StoredRunStatus,
        WorkflowRunRecord,
    )

    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        sched, store, runs = _harness(tmp_path, ownership, script={"*": "hang"})
        intended = ts(2026, 9, 8, 12, 0)
        _due(sched, store, intended)
        now = _dt.now(UTC)
        runs.save_run(
            WorkflowRunRecord(
                id="run-999009",
                status=StoredRunStatus.ADMITTED,
                resume_readiness=ResumeReadiness.NOT_APPLICABLE,
                environment=PinnedRunEnvironment(
                    deployment=_deployment(),
                    root_artifact=_artifact(),
                    child_artifacts=[],
                ),
                latest_checkpoint_id=None,
                created_at=now,
                updated_at=now,
            )
        )
        runs.mark_pending_dispatch("run-999009")
        sched.poll(intended)
        assert runs.get_run("run-999009").status.value == "failed"
        assert not runs.is_pending_dispatch("run-999009")
    finally:
        ownership.release()


def test_executing_admitted_is_abandoned_by_recovery(tmp_path: Path) -> None:
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        sched, store, runs = _harness(tmp_path, ownership, script={"*": "hang"})
        intended = ts(2026, 9, 8, 12, 0)
        _due(sched, store, intended)
        sched.poll(intended)
        run_id = runs.list_runs()[0].id
        assert runs.is_executing(run_id)
        # Fresh objects simulate a restart: executing work is abandoned, never
        # redispatched.
        fresh_runs = FileRunStore(tmp_path / "runs")
        fresh_sched_store = FileScheduleStore(tmp_path / "sched")
        diags = sched_recovery.recover(
            schedule_store=fresh_sched_store,
            run_store=fresh_runs,
            now=intended + timedelta(minutes=1),
            ownership=ownership,
        )
        assert any("failed-closed" in d for d in diags)
        assert fresh_runs.get_run(run_id).status.value == "failed"
        # A second poll after recovery never re-executes the abandoned run.
        sched2, _, runs2 = _harness(tmp_path, ownership, script={"*": "hang"})
        sched2.poll(intended + timedelta(minutes=2))
        assert runs2.get_run(run_id).status.value == "failed"
    finally:
        ownership.release()


def test_settle_hanging_run_persists_and_clears(tmp_path: Path) -> None:
    ownership = SchedulerOwnership(tmp_path, owner="test").acquire()
    try:
        sched, store, runs = _harness(tmp_path, ownership, script={"*": "hang"})
        intended = ts(2026, 9, 8, 12, 0)
        _due(sched, store, intended)
        sched.poll(intended)
        run_id = runs.list_runs()[0].id
        assert runs.is_executing(run_id)
        # The late stopped result settles without re-invoking the dispatcher.
        dispatcher = cast(SD, sched.dispatcher)
        admission = runs.get_admission(run_id)
        sched.record_stopped_execution(
            run_id, dispatcher.finish(admission, "complete"), intended
        )
        assert runs.get_run(run_id).status.value == "completed"
        assert not runs.is_executing(run_id)
        assert not runs.is_pending_dispatch(run_id)
        checkpoint = runs.get_latest_checkpoint(run_id)
        assert checkpoint.reason.value == "completed"
    finally:
        ownership.release()


def _child_poll_and_die(
    root: str, side_effect: str, intended_iso: str, sched_id: str = "a"
) -> None:
    """Child entry: poll once; the dispatcher records a side effect then dies."""
    import os as _os
    from datetime import datetime as _dt

    from tests.scheduling.controlled import (
        DictDeployments as _Dict,
    )
    from tests.scheduling.controlled import (
        ScriptedDispatcher as _SD,
    )
    from tests.scheduling.controlled import (
        fixture_environment as _env,
    )
    from wf_artifacts.runs.store import FileRunStore as _Runs
    from wf_scheduling.calendar import OneShotSource as _OneShot
    from wf_scheduling.ownership import SchedulerOwnership as _Own
    from wf_scheduling.poll import Scheduler as _Sched
    from wf_scheduling.prepare import SchedulePreparer as _Prep
    from wf_scheduling.store import FileScheduleStore as _SchedStore

    base = Path(root)
    intended = _dt.fromisoformat(intended_iso)

    def killer(admission: Any, now: _dt) -> Any:
        with open(side_effect, "a", encoding="utf-8") as fh:
            fh.write(admission.id + "\n")
        _os._exit(1)

    sched = _Sched(
        schedule_store=_SchedStore(base / "sched"),
        run_store=_Runs(base / "runs"),
        sources={sched_id: _OneShot(intended)},
        capacity=4,
        preparer=_Prep(_Dict({"dep-1": {"rev": 1, "required": []}}), _env),
        dispatcher=_SD({"*": killer}),
        ownership=_Own(base, owner="child").acquire(),
    )
    sched.poll(intended)


def test_subprocess_death_after_side_effect_dispatches_once(tmp_path: Path) -> None:
    ctx = mp.get_context("spawn")
    root = tmp_path / "death"
    root.mkdir()
    sched_store = FileScheduleStore(root / "sched")
    intended = ts(2026, 9, 8, 12, 0)
    sched_store.create_schedule(_sched_model("a"))
    sched_store.save_consumed("a", intended - timedelta(hours=1))
    side_effect = root / "effects.log"
    proc = ctx.Process(
        target=_child_poll_and_die,
        args=(str(root), str(side_effect), intended.isoformat()),
    )
    proc.start()
    proc.join(120)
    assert proc.exitcode not in (None, 0), f"child must die, got {proc.exitcode}"
    assert side_effect.read_text(encoding="utf-8").strip() != ""
    first_effects = side_effect.read_text(encoding="utf-8").strip().splitlines()

    # Restart with fresh objects: the side effect count stays one and the run
    # is abandoned, never redispatched.
    ownership = SchedulerOwnership(root, owner="parent").acquire()
    try:
        sched, store, runs = _harness(root, ownership, script={"*": "hang"})
        sched.sources["a"] = OneShotSource(intended)
        diags = sched_recovery.recover(
            schedule_store=store,
            run_store=runs,
            now=intended + timedelta(minutes=1),
            ownership=ownership,
        )
        assert any("failed-closed" in d for d in diags)
        sched.poll(intended + timedelta(minutes=1))
        effects = side_effect.read_text(encoding="utf-8").strip().splitlines()
        assert effects == first_effects
        assert len(effects) == 1
        assert runs.get_run(first_effects[0]).status.value == "failed"
    finally:
        ownership.release()