from django.db import models


class LabJob(models.Model):
    class StatusChoices(models.TextChoices):
        PENDING = "pending"
        RUNNING = "running"
        SUCCEEDED = "succeeded"
        FAILED = "failed"

    class TypeChoices(models.TextChoices):
        CREATE = "create"
        TERMINATE = "terminate"

    id = models.AutoField(primary_key=True)
    lab_instance = models.ForeignKey("labs.LabInstance", on_delete=models.CASCADE)
    type = models.CharField(max_length=20, choices=TypeChoices.choices)
    status = models.CharField(
        max_length=20, choices=StatusChoices.choices, default=StatusChoices.PENDING
    )
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        db_table = "lab_jobs"


class LabTask(models.Model):
    class StatusChoices(models.TextChoices):
        PENDING = "pending"
        SCHEDULED = "scheduled"
        RUNNING = "running"
        SUCCEEDED = "succeeded"
        FAILED = "failed"

    id = models.AutoField(primary_key=True)
    lab_instance = models.ForeignKey("labs.LabInstance", on_delete=models.CASCADE)
    lab_job = models.ForeignKey("labs.LabJob", on_delete=models.CASCADE)
    name = models.CharField(max_length=100)
    status = models.CharField(
        max_length=20, choices=StatusChoices.choices, default=StatusChoices.PENDING
    )
    payload = models.JSONField()
    retries = models.IntegerField(default=0)
    comment = models.TextField(null=True, blank=True)
    started_at = models.DateTimeField(null=True, blank=True)
    finished_at = models.DateTimeField(null=True, blank=True)
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    class Meta:
        db_table = "lab_tasks"


class LabTaskDependency(models.Model):
    id = models.AutoField(primary_key=True)
    lab_job = models.ForeignKey("labs.LabJob", on_delete=models.CASCADE)
    lab_task = models.ForeignKey(
        "labs.LabTask",
        on_delete=models.CASCADE,
        related_name="lab_task_dependencies",
    )
    depends_on = models.ForeignKey(
        "labs.LabTask",
        on_delete=models.CASCADE,
        related_name="lab_task_dependents",
    )

    class Meta:
        db_table = "lab_task_dependencies"