activity-core/src/activity_core/schedule_manager.py

"""Temporal Schedule management for activity-core.

T22: upsert_schedule, delete_schedule, list_schedules
T24: misfire_policy → ScheduleOverlapPolicy mapping (all three policies)

Schedule ID convention: activity-schedule-{activity_definition.id}
Workflow triggered: RunActivityWorkflow on orchestrator-tq
"""

from __future__ import annotations

from datetime import datetime, timedelta, timezone
from uuid import UUID

from temporalio.client import (
    Client,
    Schedule,
    ScheduleActionStartWorkflow,
    ScheduleAlreadyRunningError,
    ScheduleBackfill,
    ScheduleCalendarSpec,
    ScheduleHandle,
    ScheduleOverlapPolicy,
    SchedulePolicy,
    ScheduleRange,
    ScheduleSpec,
    ScheduleState,
    ScheduleUpdate,
    ScheduleUpdateInput,
)
from temporalio.service import RPCError

from activity_core.models import ActivityDefinition, CronTriggerConfig, ScheduledTriggerConfig

_ORCHESTRATOR_TASK_QUEUE = "orchestrator-tq"

# Trigger_key sentinel used when a workflow is started by a Temporal Schedule.
# RunActivityWorkflow detects this value and derives run dedup key from workflow_id.
SCHEDULED_TRIGGER_KEY = "scheduled"

# T24: misfire_policy → ScheduleOverlapPolicy
_MISFIRE_TO_OVERLAP: dict[str, ScheduleOverlapPolicy] = {
    "skip": ScheduleOverlapPolicy.SKIP,
    "catchup": ScheduleOverlapPolicy.BUFFER_ALL,
    "compress": ScheduleOverlapPolicy.BUFFER_ONE,
}


def schedule_id(activity_id: str | UUID) -> str:
    """Return the canonical Temporal Schedule ID for an ActivityDefinition."""
    return f"activity-schedule-{activity_id}"


def smoke_schedule_id(activity_id: str | UUID) -> str:
    """Return the one-shot smoke-test Schedule ID for an ActivityDefinition."""
    return f"activity-smoke-test-{activity_id}"


def _overlap_policy(misfire_policy: str) -> ScheduleOverlapPolicy:
    return _MISFIRE_TO_OVERLAP.get(misfire_policy, ScheduleOverlapPolicy.SKIP)


def _build_schedule(defn: ActivityDefinition) -> Schedule:
    """Construct a Temporal Schedule object from a cron ActivityDefinition."""
    assert isinstance(defn.trigger_config, CronTriggerConfig)
    cfg: CronTriggerConfig = defn.trigger_config

    # Workflow ID uses ${firstScheduledTime} so each schedule fire gets a
    # unique workflow ID, enabling replay/audit without ID conflicts.
    action = ScheduleActionStartWorkflow(
        "RunActivityWorkflow",
        args=[str(defn.id), SCHEDULED_TRIGGER_KEY, None],
        id=f"activity-{defn.id}:${{firstScheduledTime}}",
        task_queue=_ORCHESTRATOR_TASK_QUEUE,
    )

    spec = ScheduleSpec(
        cron_expressions=[cfg.cron_expression],
        time_zone_name=cfg.timezone,
        jitter=timedelta(seconds=cfg.jitter_seconds) if cfg.jitter_seconds else None,
    )

    policy = SchedulePolicy(overlap=_overlap_policy(cfg.misfire_policy))
    state = ScheduleState(paused=not defn.enabled)

    return Schedule(action=action, spec=spec, policy=policy, state=state)


def _onetime_schedule_id(activity_id: str | UUID) -> str:
    return f"activity-schedule-{activity_id}-once"


def _build_onetime_schedule(defn: ActivityDefinition) -> tuple[str, Schedule]:
    """Build a one-off Temporal Schedule that fires once at defn.trigger_config.at.

    Returns (schedule_id, Schedule).
    Uses ScheduleState(remaining_actions=1) so the schedule self-disarms after firing.
    """
    assert isinstance(defn.trigger_config, ScheduledTriggerConfig)
    cfg: ScheduledTriggerConfig = defn.trigger_config
    at = cfg.at

    action = ScheduleActionStartWorkflow(
        "RunActivityWorkflow",
        args=[str(defn.id), SCHEDULED_TRIGGER_KEY, at.isoformat(), None],
        id=f"activity-{defn.id}:once",
        task_queue=_ORCHESTRATOR_TASK_QUEUE,
    )

    # Calendar spec pinned to the exact minute — combined with remaining_actions=1
    # this fires exactly once at the specified time.
    spec = ScheduleSpec(
        calendars=[
            ScheduleCalendarSpec(
                second=[ScheduleRange(0)],
                minute=[ScheduleRange(at.minute)],
                hour=[ScheduleRange(at.hour)],
                day_of_month=[ScheduleRange(at.day)],
                month=[ScheduleRange(at.month)],
                year=[ScheduleRange(at.year)],
            )
        ],
        time_zone_name=cfg.timezone,
    )

    state = ScheduleState(
        limited_actions=True,
        remaining_actions=1,
        paused=not defn.enabled,
    )

    sid = _onetime_schedule_id(defn.id)
    return sid, Schedule(action=action, spec=spec, state=state)


def _build_smoke_test_schedule(
    defn: ActivityDefinition,
    fire_at: datetime,
) -> tuple[str, str, Schedule]:
    """Build a one-shot smoke Schedule for an enabled cron ActivityDefinition."""
    if not isinstance(defn.trigger_config, CronTriggerConfig):
        raise ValueError("schedule smoke tests require trigger_type='cron'")
    if not defn.enabled:
        raise ValueError("schedule smoke tests require an enabled ActivityDefinition")

    at = fire_at.astimezone(timezone.utc)
    token = at.strftime("%Y%m%dT%H%M%SZ")
    workflow_id_prefix = f"activity-{defn.id}:smoke-{token}"
    trigger_key = f"schedule-smoke-{token}"

    action = ScheduleActionStartWorkflow(
        "RunActivityWorkflow",
        args=[str(defn.id), trigger_key, at.isoformat(), None],
        id=workflow_id_prefix,
        task_queue=_ORCHESTRATOR_TASK_QUEUE,
    )

    spec = ScheduleSpec(
        calendars=[
            ScheduleCalendarSpec(
                second=[ScheduleRange(at.second)],
                minute=[ScheduleRange(at.minute)],
                hour=[ScheduleRange(at.hour)],
                day_of_month=[ScheduleRange(at.day)],
                month=[ScheduleRange(at.month)],
                year=[ScheduleRange(at.year)],
            )
        ],
        time_zone_name="UTC",
    )

    state = ScheduleState(
        limited_actions=True,
        remaining_actions=1,
        paused=False,
    )

    return (
        smoke_schedule_id(defn.id),
        workflow_id_prefix,
        Schedule(action=action, spec=spec, state=state),
    )


async def cancel_scheduled(client: Client, activity_id: str | UUID) -> None:
    """Delete the one-off Temporal Schedule for a ScheduledTriggerConfig definition.

    No-op if the schedule does not exist.
    """
    handle = client.get_schedule_handle(_onetime_schedule_id(activity_id))
    try:
        await handle.delete()
    except RPCError:
        pass


async def schedule_smoke_test(
    client: Client,
    defn: ActivityDefinition,
    *,
    delay: timedelta = timedelta(minutes=1),
    now: datetime | None = None,
) -> tuple[str, str, datetime]:
    """Schedule a one-shot smoke run for a recurring ActivityDefinition.

    Returns ``(schedule_id, workflow_id_prefix, fire_at)``. Temporal appends
    the scheduled fire time to workflow IDs created by schedules.
    """
    base = now or datetime.now(tz=timezone.utc)
    if base.tzinfo is None:
        base = base.replace(tzinfo=timezone.utc)
    fire_at = (base + delay).astimezone(timezone.utc)
    sid, workflow_id_prefix, sched = _build_smoke_test_schedule(defn, fire_at)
    try:
        await client.create_schedule(sid, sched)
    except (RPCError, ScheduleAlreadyRunningError):
        handle = client.get_schedule_handle(sid)

        async def _updater_smoke(inp: ScheduleUpdateInput) -> ScheduleUpdate:  # noqa: ARG001
            return ScheduleUpdate(schedule=sched)

        await handle.update(_updater_smoke)
        await handle.unpause()
    return sid, workflow_id_prefix, fire_at


async def delete_smoke_test_schedule(client: Client, activity_id: str | UUID) -> None:
    """Delete the smoke-test Schedule for the given activity_id if present."""
    handle = client.get_schedule_handle(smoke_schedule_id(activity_id))
    try:
        await handle.delete()
    except RPCError:
        pass


async def upsert_schedule(client: Client, defn: ActivityDefinition) -> ScheduleHandle:
    """Create or update a Temporal Schedule for a cron or scheduled ActivityDefinition.

    - For cron: creates/updates the recurring schedule.
    - For scheduled: creates a one-off schedule (remaining_actions=1).
    - If enabled=False the schedule is created paused.
    - For cron with misfire_policy='catchup', triggers a backfill covering the
      last hour after each upsert to replay any recently missed fires.

    Returns the ScheduleHandle for the created/updated schedule.
    """
    if isinstance(defn.trigger_config, ScheduledTriggerConfig):
        sid, sched = _build_onetime_schedule(defn)
        try:
            handle = await client.create_schedule(sid, sched)
        except RPCError:
            handle = client.get_schedule_handle(sid)
            async def _updater_once(inp: ScheduleUpdateInput) -> ScheduleUpdate:  # noqa: ARG001
                return ScheduleUpdate(schedule=sched)
            await handle.update(_updater_once)
        return handle

    if not isinstance(defn.trigger_config, CronTriggerConfig):
        raise ValueError(
            f"upsert_schedule requires trigger_type='cron' or 'scheduled', "
            f"got {defn.trigger_config.trigger_type!r}"
        )

    sid = schedule_id(defn.id)
    sched = _build_schedule(defn)

    try:
        handle = await client.create_schedule(sid, sched)
    except (RPCError, ScheduleAlreadyRunningError):
        # Schedule already exists — update it in place.
        handle = client.get_schedule_handle(sid)

        async def _updater(input: ScheduleUpdateInput) -> ScheduleUpdate:  # noqa: ARG001
            return ScheduleUpdate(schedule=sched)

        await handle.update(_updater)

        # Sync pause state explicitly (update replaces the schedule object
        # but pause state is part of ScheduleState, already embedded above).
        if defn.enabled:
            await handle.unpause()
        else:
            await handle.pause(note="disabled via upsert_schedule")

    # T24 catchup: backfill any fires missed in the last hour.
    if isinstance(defn.trigger_config, CronTriggerConfig):
        if defn.trigger_config.misfire_policy == "catchup":
            now = datetime.now(tz=timezone.utc)
            backfill_start = now - timedelta(hours=1)
            await handle.backfill(
                ScheduleBackfill(
                    start_at=backfill_start,
                    end_at=now,
                    overlap=ScheduleOverlapPolicy.BUFFER_ALL,
                )
            )

    return handle


async def delete_schedule(client: Client, activity_id: str | UUID) -> None:
    """Delete the Temporal Schedule for the given activity_id.

    No-op if the schedule does not exist.
    """
    handle = client.get_schedule_handle(schedule_id(activity_id))
    try:
        await handle.delete()
    except RPCError:
        pass  # Not found — treat as success.


async def list_schedules(client: Client) -> list[dict]:
    """Enumerate all activity-core Temporal Schedules.

    Returns a list of dicts: [{"schedule_id": str, "activity_id": str}, ...]
    """
    prefix = "activity-schedule-"
    results: list[dict] = []
    async for entry in await client.list_schedules():
        if entry.id.startswith(prefix):
            results.append(
                {
                    "schedule_id": entry.id,
                    "activity_id": entry.id[len(prefix) :],
                }
            )
    return results