import sky


class SkypilotJob(object):
    def __init__(self, cluster_name, job_id):
        self.cluster_name = cluster_name
        self.job_id = job_id

    def __repr__(self):
        return "{}('{}':'{}')".format(self.__class__.__name__, self.cluster_name, self.job_id)

    @property
    def id(self):
        return self.job_id

    @property
    def status(self) -> sky.JobStatus:
        # sky.job_status returns request id which then returns {job_id: JobStatus}
        request_id = sky.job_status(self.cluster_name, [self.job_id])
        status_dict = sky.get(request_id)
        return status_dict.get(self.job_id)

    @property
    def is_waiting(self):
        s = self.status
        return s is not None and s in [sky.JobStatus.INIT, sky.JobStatus.PENDING]

    @property
    def is_running(self):
        s = self.status
        return s is not None and not s.is_terminal()

    @property
    def has_failed(self):
        return self.status in [sky.JobStatus.FAILED, sky.JobStatus.FAILED_SETUP, sky.JobStatus.FAILED_DRIVER]

    @property
    def has_succeeded(self):
        return self.status in [sky.JobStatus.SUCCEEDED]

    @property
    def has_finished(self):
        return self.has_succeeded or self.has_failed or self.status in [sky.JobStatus.CANCELLED]

    @property
    def message(self):
        return self.status.value if self.status else "Unknown"

    def kill(self):
        if not self.has_finished:
            sky.cancel(self.cluster_name, job_ids=[self.job_id])
