import traceback
import logging

from celery import shared_task
from django.utils import timezone
from django.utils.module_loading import import_string
from django.db import connection, transaction, OperationalError
from datetime import timedelta

from labs.models import (
    LabInstance as DBLabInstance,
    LabTask as DBLabTask,
    LabJob as DBLabJob,
)
from labs.lab_manager_v2.lab_instance import LabInstance
from labs.lab_manager_v2.exceptions import LabTaskRecoverableError
from labs.lab_manager_v2.tasks.base import BaseLabTask

logger = logging.getLogger(__name__)


@shared_task(bind=True, max_retries=100)
def run_lab_task_v2(self, lab_task_id: int):
    try:
        started_at = timezone.now()
        db_lab_task: DBLabTask = DBLabTask.objects.select_related(
            "lab_job", "lab_instance"
        ).get(id=lab_task_id)
        logger.info(
            f"Running lab task {db_lab_task.name} for lab {db_lab_task.lab_instance.unique_tag}"
        )
        db_lab_instance: DBLabInstance = db_lab_task.lab_instance
        db_lab_job: DBLabJob = db_lab_task.lab_job

        if db_lab_task.status in (
            DBLabTask.StatusChoices.SUCCEEDED,
            DBLabTask.StatusChoices.FAILED,
        ):
            logger.error(
                f"Lab task {lab_task_id} is already succeeded or failed, skipping"
            )
            return

        lab_instance = LabInstance(db_lab_instance)
        job_manager = lab_instance.get_job_manager(db_lab_job)
        task_class = import_string(db_lab_task.name)
        task: BaseLabTask = task_class(db_lab_task)

        db_lab_task.status = DBLabTask.StatusChoices.RUNNING
        db_lab_task.started_at = started_at
        job_manager.save_task(db_lab_task)

        task.run()

        # Make sure the connection is available
        # Because we are using TiDB Serverless, it's connection is not stable after a long time
        connection.close_if_unusable_or_obsolete()
        db_lab_task.finished_at = timezone.now()
        db_lab_task.status = DBLabTask.StatusChoices.SUCCEEDED
        job_manager.save_task(db_lab_task)
        job_manager.schedule_next_task()
    except OperationalError as e:
        logger.error(f"Error getting lab task {lab_task_id}: {e}")
        raise self.retry(countdown=10)
    except LabTaskRecoverableError as e:
        if db_lab_task.retries < task.max_retries:
            logger.info(
                f"Recoverable task {db_lab_task.name} for lab {db_lab_instance.unique_tag} failed: {e}, retrying..."
            )
            db_lab_task.retries += 1
            job_manager.save_task(db_lab_task)
            self.retry(countdown=task.retry_interval_seconds)
        else:
            logger.error(
                f"Recoverable task {db_lab_task.name} for lab {db_lab_instance.unique_tag} failed: {e}, giving up"
            )
            db_lab_task.comment = str(e)
            db_lab_task.status = DBLabTask.StatusChoices.FAILED
            job_manager.save_task(db_lab_task)
            return
    except Exception as e:
        logger.error(
            f"Unrecoverable task {db_lab_task.name} for lab {db_lab_instance.unique_tag} failed: {traceback.format_exc()}"
        )
        db_lab_task.comment = str(e)
        db_lab_task.status = DBLabTask.StatusChoices.FAILED
        job_manager.save_task(db_lab_task)
        return


@shared_task
def recycle_timeout_lab_instances():
    with transaction.atomic():
        db_lab_instances = DBLabInstance.objects.select_for_update().filter(
            status=DBLabInstance.LabInstanceStatus.RUNNING,
            scheduled_terminate_at__lt=timezone.now(),
        )
        for db_lab_instance in db_lab_instances:
            lab_instance = LabInstance(db_lab_instance)
            logger.info(f"Recycling timeout lab instance {db_lab_instance.unique_tag}")
            lab_instance.terminate_lab()
        logger.info(f"Recycled {len(db_lab_instances)} timeout lab instances")


DEFAULT_EXIT_LAB_INSTANCE_TIMEOUT_DAYS = 7


@shared_task
def recycle_stale_lab_instances(
    timeout_days: int = DEFAULT_EXIT_LAB_INSTANCE_TIMEOUT_DAYS,
):
    # If the lab instance is terminated for more than 7 days,
    # user may already forget about the lab instance,
    # so we need to exit the lab instance for them.
    # This task can help us to free up the resources.
    DBLabInstance.objects.filter(
        status=DBLabInstance.LabInstanceStatus.TERMINATED,
        exited_at__isnull=True,
        terminated_at__lt=timezone.now() - timedelta(days=timeout_days),
    ).update(exited_at=timezone.now())