"""Execution control helpers for generation harness workflows."""

from typing import TYPE_CHECKING

from sqlalchemy.ext.asyncio import AsyncSession

from app.services.generation_jobs import finish_generation_job

if TYPE_CHECKING:
    from app.db.models import GenerationJob, Story


class GenerationJobCanceledError(Exception):
    """Raised when a running worker job has been canceled by the user."""


class ExecutionControl:
    """Runtime control surface for cancelable generation workflows."""

    def __init__(self, db: AsyncSession):
        self.db = db

    async def stop_if_cancel_requested(
        self,
        *,
        job: "GenerationJob | None",
        story: "Story | None" = None,
    ) -> None:
        """Stop a worker-owned job at the next safe checkpoint after cancellation."""

        if job is None:
            return

        await self.db.refresh(job)
        if job.current_step != "cancel_requested":
            return

        await finish_generation_job(
            self.db,
            job=job,
            story=story,
            status="canceled",
            current_step="generation_canceled",
            error_message="Generation canceled by user.",
            message="Generation job was canceled after a user request.",
        )
        raise GenerationJobCanceledError()