from dataclasses import dataclass
from enum import Enum
from time import monotonic


class JobStatus(str, Enum):
    CREATED = "created"
    BLOCKED = "blocked"
    ROUTED = "routed"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    TIMED_OUT = "timed_out"


ALLOWED_TRANSITIONS = {
    JobStatus.CREATED: {
        JobStatus.BLOCKED,
        JobStatus.ROUTED,
    },
    JobStatus.ROUTED: {
        JobStatus.RUNNING,
        JobStatus.FAILED,
    },
    JobStatus.RUNNING: {
        JobStatus.COMPLETED,
        JobStatus.FAILED,
        JobStatus.TIMED_OUT,
    },
}


class InvalidTransition(ValueError):
    pass


@dataclass
class JobLifecycle:
    job_id: str
    status: JobStatus = JobStatus.CREATED
    started_at: float | None = None
    finished_at: float | None = None

    def transition(self, new_status: JobStatus):
        allowed = ALLOWED_TRANSITIONS.get(
            self.status, set()
        )

        if new_status not in allowed:
            raise InvalidTransition(
                f"{self.status.value} -> "
                f"{new_status.value} not allowed"
            )

        self.status = new_status

        if new_status == JobStatus.RUNNING:
            self.started_at = monotonic()

        if new_status in {
            JobStatus.COMPLETED,
            JobStatus.FAILED,
            JobStatus.TIMED_OUT,
            JobStatus.BLOCKED,
        }:
            self.finished_at = monotonic()