import logging

from django.db import transaction
from django.db.models import Exists, OuterRef, Count
from django.utils import timezone

from labs.models import (
    LabJob as DBLabJob,
    LabTask as DBLabTask,
    LabTaskDependency as LabTaskDependency,
    LabInstance as DBLabInstance,
)
from labs.lab_manager_v2.lab_spec import (
    LabSpec,
    CloudStackRequirement,
    TiDBServerlessClusterRequirement,
    AWSAuroraServerlessClusterRequirement,
)
from labs.billing.lab import refund_credits_for_lab_failure

logger = logging.getLogger(__name__)


class CeleryLabJob:
    def __init__(self, db_lab_job: DBLabJob):
        self.db_lab_job = db_lab_job
        self.db_lab_instance = db_lab_job.lab_instance

    @classmethod
    def create(
        cls, db_lab_instance: DBLabInstance, lab_spec: LabSpec
    ) -> "CeleryLabJob":
        """
        Creates a new lab job and related sub-tasks, task dependencies.
        This is used to create a new lab job for a lab instance.
        """
        raise NotImplementedError

    def can_schedule_task(self, db_lab_instance: DBLabInstance) -> bool:
        """
        Returns True if the lab instance can schedule a task.
        This is used to check if the lab instance is in the correct status to schedule a task.
        """
        raise NotImplementedError

    @transaction.atomic
    def schedule_next_task(self) -> None:
        logger.info(f"Scheduling next task for lab {self.db_lab_instance.id}")
        db_lab_instance = DBLabInstance.objects.select_for_update().get(
            id=self.db_lab_instance.id
        )
        if not self.can_schedule_task(db_lab_instance):
            return

        runnable_tasks = self.find_runnable_tasks(self.db_lab_job)
        from labs.tasks.lab_v2 import run_lab_task_v2

        for task in runnable_tasks:
            task.status = DBLabTask.StatusChoices.SCHEDULED
            task.save()
            run_lab_task_v2.delay(task.id)

        logger.info(
            f"Scheduled {len(runnable_tasks)} tasks for lab {self.db_lab_instance.id}"
        )

    def find_runnable_tasks(self, lab_job: DBLabJob):
        """
        Finds tasks for a given lab_job that have no dependencies or whose dependencies are all completed.
        Only tasks with 'pending' status are considered.
        """
        # Subquery to find any dependency task that is not yet 'succeeded'.
        # For each LabTask in the outer query, this subquery will check its dependencies.
        uncompleted_dependencies = LabTaskDependency.objects.filter(
            # This links the subquery to each task in the main query
            lab_task=OuterRef("pk"),
        ).exclude(
            # We are only interested in dependencies that are NOT succeeded.
            # If all dependencies are succeeded, this subquery will return an empty set for the task.
            depends_on__status=DBLabTask.StatusChoices.SUCCEEDED
        )

        # The main query for LabTask
        runnable_tasks = (
            DBLabTask.objects.filter(
                lab_job=lab_job, status=DBLabTask.StatusChoices.PENDING
            )
            .annotate(
                # For each task, check if there exists any uncompleted dependency.
                # The result is stored in the 'has_uncompleted_deps' annotation.
                has_uncompleted_deps=Exists(uncompleted_dependencies)
            )
            .filter(
                # We only want tasks that do NOT have any uncompleted dependencies.
                # This includes tasks with no dependencies at all, or tasks where all dependencies are 'succeeded'.
                has_uncompleted_deps=False
            )
        )
        return runnable_tasks

    def save_task(self, lab_task: DBLabTask):
        """
        Saves the status of a task, update the status of the lab instance if necessary.
        """
        raise NotImplementedError


class CeleryLabCreateJob(CeleryLabJob):
    @classmethod
    def create(
        cls, db_lab_instance: DBLabInstance, lab_spec: LabSpec
    ) -> "CeleryLabCreateJob":
        from labs.lab_manager_v2.tasks import (
            AWSCloudFormationStackCreateTask,
            AWSCloudFormationStackReadyCheckTask,
            TagAWSEC2InstanceTask,
            AlicloudROSStackCreateTask,
            AlicloudROSStackReadyCheckTask,
            AlicloudTagECSInstanceTask,
            TiDBCloudServerlessClusterCreateTask,
            ExecuteBootstrapScriptTask,
            SetupHTTPSReverseProxyTask,
            EC2ReadyCheckTask,
            AWSAuroraServerlessClusterCreateTask,
            AWSAuroraServerlessClusterReadyCheckTask,
        )
        db_lab_job = DBLabJob.objects.create(
            lab_instance=db_lab_instance, type=DBLabJob.TypeChoices.CREATE
        )
        ec2_depends_on = []
        for req in lab_spec.requirements:
            if isinstance(req, CloudStackRequirement):
                if db_lab_instance.cloud_provider == "aws":
                    cfn_task = AWSCloudFormationStackCreateTask.create(
                        db_lab_job=db_lab_job
                    )
                    cfn_ready_task = AWSCloudFormationStackReadyCheckTask.create(
                        db_lab_job=db_lab_job, depends_on=[cfn_task]
                    )
                    tag_ec2_task = TagAWSEC2InstanceTask.create(
                        db_lab_job=db_lab_job, depends_on=[cfn_ready_task]
                    )
                    ec2_depends_on.append(tag_ec2_task)
                elif db_lab_instance.cloud_provider == "alicloud":
                    ros_task = AlicloudROSStackCreateTask.create(db_lab_job=db_lab_job)
                    ros_ready_task = AlicloudROSStackReadyCheckTask.create(
                        db_lab_job=db_lab_job, depends_on=[ros_task]
                    )
                    tag_ecs_task = AlicloudTagECSInstanceTask.create(
                        db_lab_job=db_lab_job, depends_on=[ros_ready_task]
                    )
                    ec2_depends_on.append(tag_ecs_task)
                else:
                    raise ValueError(
                        f"Unsupported cloud provider: {db_lab_instance.cloud_provider}"
                    )
            elif isinstance(req, TiDBServerlessClusterRequirement):
                tidb_serverless_task = TiDBCloudServerlessClusterCreateTask.create(
                    db_lab_job=db_lab_job
                )
                ec2_depends_on.append(tidb_serverless_task)
            elif isinstance(req, AWSAuroraServerlessClusterRequirement):
                aws_aurora_serverless_task = AWSAuroraServerlessClusterCreateTask.create(
                    db_lab_job=db_lab_job
                )
                aws_aurora_serverless_ready_task = AWSAuroraServerlessClusterReadyCheckTask.create(
                    db_lab_job=db_lab_job, depends_on=[aws_aurora_serverless_task]
                )
                ec2_depends_on.append(aws_aurora_serverless_ready_task)
            else:
                raise ValueError(f"Unsupported requirement: {req}")

        ec2_ready_task = EC2ReadyCheckTask.create(
            db_lab_job=db_lab_job, depends_on=ec2_depends_on
        )

        if any(n.https_reverse_proxy for n in lab_spec.topology):
            SetupHTTPSReverseProxyTask.create(
                db_lab_job=db_lab_job, depends_on=[ec2_ready_task]
            )

        bootstrap_tasks = []
        for index, step in enumerate(lab_spec.bootstrap):
            if index == 0:
                depends_on = [ec2_ready_task]
            else:
                depends_on = [bootstrap_tasks[index - 1]]
            bootstrap_task = ExecuteBootstrapScriptTask.create(
                db_lab_job=db_lab_job, depends_on=depends_on, payload={"index": index}
            )
            bootstrap_tasks.append(bootstrap_task)
        # TODO: check if there is a circular dependency in the tasks
        return cls(db_lab_job=db_lab_job)

    def can_schedule_task(self, db_lab_instance: DBLabInstance) -> bool:
        can = db_lab_instance.status == DBLabInstance.LabInstanceStatus.CREATING
        if not can:
            logger.info(
                f"Lab {self.db_lab_instance.id} is not in {DBLabInstance.LabInstanceStatus.CREATING} status, skipping"
            )
        return can

    @transaction.atomic
    def save_task(self, db_lab_task: DBLabTask):
        db_lab_instance = DBLabInstance.objects.select_for_update().get(
            id=self.db_lab_instance.id
        )
        db_lab_task.save()
        task_status_map = dict(
            DBLabTask.objects.filter(lab_job=self.db_lab_job)
            .values("status")
            .annotate(count=Count("status"))
            .values_list("status", "count")
        )
        task_count = sum(task_status_map.values())
        if task_status_map.get(DBLabTask.StatusChoices.FAILED, 0) > 0:
            self.db_lab_job.status = DBLabJob.StatusChoices.FAILED
            self.db_lab_job.save()
            db_lab_instance.status = DBLabInstance.LabInstanceStatus.CREATING_FAILED
            db_lab_instance.save(update_fields=["status"])
            refund_credits_for_lab_failure(lab_instance=db_lab_instance)
        elif task_status_map.get(DBLabTask.StatusChoices.SUCCEEDED, 0) == task_count:
            self.db_lab_job.status = DBLabJob.StatusChoices.SUCCEEDED
            self.db_lab_job.save()
            db_lab_instance.status = DBLabInstance.LabInstanceStatus.RUNNING
            db_lab_instance.ready_at = timezone.now()
            db_lab_instance.save(update_fields=["status", "ready_at"])


class CeleryLabTerminateJob(CeleryLabJob):
    @classmethod
    def create(
        cls, db_lab_instance: DBLabInstance, lab_spec: LabSpec
    ) -> "CeleryLabTerminateJob":
        from labs.lab_manager_v2.tasks import (
            AWSCloudFormationStackDeleteTask,
            TiDBCloudServerlessClusterDeleteTask,
            AlicloudROSStackDeleteTask,
            AWSAuroraServerlessClusterDeleteTask,
        )
        db_lab_job = DBLabJob.objects.create(
            lab_instance=db_lab_instance, type=DBLabJob.TypeChoices.TERMINATE
        )
        for req in lab_spec.requirements:
            if isinstance(req, CloudStackRequirement):
                if db_lab_instance.cloud_provider == "aws":
                    AWSCloudFormationStackDeleteTask.create(db_lab_job=db_lab_job)
                elif db_lab_instance.cloud_provider == "alicloud":
                    AlicloudROSStackDeleteTask.create(db_lab_job=db_lab_job)
                else:
                    raise ValueError(
                        f"Unsupported cloud provider: {db_lab_instance.cloud_provider}"
                    )
            elif isinstance(req, TiDBServerlessClusterRequirement):
                TiDBCloudServerlessClusterDeleteTask.create(db_lab_job=db_lab_job)
            elif isinstance(req, AWSAuroraServerlessClusterRequirement):
                AWSAuroraServerlessClusterDeleteTask.create(db_lab_job=db_lab_job)
            else:
                raise ValueError(f"Unsupported requirement: {req}")
        return cls(db_lab_job=db_lab_job)

    def can_schedule_task(self, db_lab_instance: DBLabInstance) -> bool:
        can = db_lab_instance.status == DBLabInstance.LabInstanceStatus.TERMINATING
        if not can:
            logger.info(
                f"Lab {self.db_lab_instance.id} is not in {DBLabInstance.LabInstanceStatus.TERMINATING} status, skipping"
            )
        return can

    @transaction.atomic
    def save_task(self, db_lab_task: DBLabTask):
        db_lab_instance = DBLabInstance.objects.select_for_update().get(
            id=self.db_lab_instance.id
        )
        db_lab_task.save()
        task_status_map = dict(
            DBLabTask.objects.filter(lab_job=self.db_lab_job)
            .values("status")
            .annotate(count=Count("status"))
            .values_list("status", "count")
        )
        task_count = sum(task_status_map.values())
        if task_status_map.get(DBLabTask.StatusChoices.FAILED, 0) > 0:
            self.db_lab_job.status = DBLabJob.StatusChoices.FAILED
            self.db_lab_job.save()
            db_lab_instance.status = DBLabInstance.LabInstanceStatus.TERMINATING_FAILED
            db_lab_instance.save(update_fields=["status"])
        elif task_status_map.get(DBLabTask.StatusChoices.SUCCEEDED, 0) == task_count:
            self.db_lab_job.status = DBLabJob.StatusChoices.SUCCEEDED
            self.db_lab_job.save()
            db_lab_instance.status = DBLabInstance.LabInstanceStatus.TERMINATED
            db_lab_instance.terminated_at = timezone.now()
            db_lab_instance.save(update_fields=["status", "terminated_at"])