from typing import Optional, List

from labs.models import (
    LabInstance as DBLabInstance,
    LabJob as DBLabJob,
    LabTask as DBLabTask,
    LabTaskDependency as DBLabTaskDependency,
)
from labs.lab_manager_v2.lab_spec import LabSpec, load_lab_spec


class BaseLabTask:
    db_task: DBLabTask
    db_lab_instance: DBLabInstance
    lab_spec: LabSpec
    max_retries: int = 3
    retry_interval_seconds: int = 10

    def __init__(self, db_lab_task: DBLabTask) -> None:
        self.db_task = db_lab_task
        self.db_lab_instance = db_lab_task.lab_instance
        self.lab_spec = load_lab_spec(self.db_lab_instance.lab_code)

    @classmethod
    def create(
        cls,
        db_lab_job: DBLabJob,
        payload: dict | None = None,
        depends_on: Optional[List[DBLabTask]] = None,
    ) -> DBLabTask:
        if payload is None:
            payload = {}
        db_task = DBLabTask.objects.create(
            lab_instance=db_lab_job.lab_instance,
            lab_job=db_lab_job,
            name=f"{cls.__module__}.{cls.__name__}",
            payload=payload,
        )
        if depends_on:
            db_task_dependencies = []
            for depend_db_task in depends_on:
                db_task_dependencies.append(
                    DBLabTaskDependency(
                        lab_job=db_lab_job,
                        lab_task=db_task,
                        depends_on=depend_db_task,
                    )
                )
            DBLabTaskDependency.objects.bulk_create(db_task_dependencies)
        return db_task

    def run(self) -> None:
        raise NotImplementedError