core/tasks.py

"""
Async tasks for PyRunner using django-q2.

This module contains task functions that are executed asynchronously
by django-q2 workers.
"""

import logging
from uuid import UUID

from django.utils import timezone
from django_q.tasks import async_task

from core.models import Run, Script, ScriptSchedule, GlobalSettings, PackageOperation

logger = logging.getLogger(__name__)


def execute_run_task(run_id: str, webhook_data: dict | None = None) -> dict:
    """
    Execute a script run asynchronously.

    This is the task function called by django-q2 workers.
    It fetches the Run by ID and delegates to the executor.

    Args:
        run_id: The UUID of the Run record (as string)
        webhook_data: Optional webhook data to inject as environment variables

    Returns:
        dict: Execution result summary for logging/monitoring
    """
    from core.executor import execute_run  # Import here to avoid circular imports
    from core.services.notification_service import NotificationService

    try:
        run = Run.objects.select_related("script", "script__environment").get(
            id=UUID(run_id)
        )
    except Run.DoesNotExist:
        logger.error(f"Run {run_id} not found - task cannot execute")
        return {
            "success": False,
            "run_id": run_id,
            "error": "Run not found",
        }
    except ValueError as e:
        logger.error(f"Invalid run_id format: {run_id} - {e}")
        return {
            "success": False,
            "run_id": run_id,
            "error": f"Invalid UUID format: {e}",
        }

    execute_run(run, webhook_data=webhook_data)
    run.refresh_from_db()

    # Send notifications after run completion
    try:
        notification_results = NotificationService.send_notification(run)
        if notification_results.get("email_sent") or notification_results.get("webhook_sent"):
            logger.info(f"Notifications sent for run {run_id}: {notification_results}")
    except Exception as e:
        logger.error(f"Failed to send notifications for run {run_id}: {e}")

    logger.info(f"Task completed for Run {run_id} with status {run.status}")

    return {
        "success": run.status == Run.Status.SUCCESS,
        "run_id": run_id,
        "status": run.status,
        "exit_code": run.exit_code,
    }


def queue_script_run(run: Run, webhook_data: dict | None = None) -> str:
    """
    Queue a Run for async execution.

    This is the main entry point for queuing script runs.
    It handles setting up the async task and storing the task_id.

    Args:
        run: The Run model instance to execute
        webhook_data: Optional webhook data to inject as environment variables

    Returns:
        str: The django-q2 task ID

    Raises:
        ValueError: If the run is not in PENDING status
    """
    if run.status != Run.Status.PENDING:
        raise ValueError(
            f"Cannot queue run {run.id}: status is {run.status}, expected PENDING"
        )

    task_id = async_task(
        "core.tasks.execute_run_task",
        str(run.id),
        webhook_data,
        task_name=f"run-{run.id}",
        timeout=run.script.timeout_seconds + 60,
    )

    run.task_id = task_id
    run.save(update_fields=["task_id"])

    logger.info(f"Queued Run {run.id} as task {task_id}")

    return task_id


def execute_scheduled_run(script_id: str) -> dict:
    """
    Execute a scheduled script run.
    Called by django-q2 scheduler.

    Args:
        script_id: The UUID of the Script (as string)

    Returns:
        dict: Execution result summary
    """
    from core.services.schedule_service import ScheduleService

    # Check global pause
    settings = GlobalSettings.get_settings()
    if settings.schedules_paused:
        logger.info(f"Scheduled run for script {script_id} skipped - schedules paused")
        return {"success": False, "error": "Schedules globally paused"}

    try:
        script = Script.objects.get(id=UUID(script_id))
    except Script.DoesNotExist:
        logger.error(f"Script {script_id} not found for scheduled run")
        return {"success": False, "error": "Script not found"}
    except ValueError as e:
        logger.error(f"Invalid script_id format: {script_id} - {e}")
        return {"success": False, "error": f"Invalid UUID format: {e}"}

    # Check if script can run (enabled and not archived)
    if not script.can_run:
        reason = "archived" if script.is_archived else "disabled"
        logger.info(f"Scheduled run for script {script.name} skipped - script {reason}")
        return {"success": False, "error": f"Script {reason}"}

    # Check if schedule is active
    try:
        schedule = script.schedule
        if not schedule.is_active:
            logger.info(
                f"Scheduled run for script {script.name} skipped - schedule inactive"
            )
            return {"success": False, "error": "Schedule inactive"}
    except ScriptSchedule.DoesNotExist:
        logger.error(f"No schedule found for script {script.name}")
        return {"success": False, "error": "No schedule"}

    # Create the run
    run = Run.objects.create(
        script=script,
        status=Run.Status.PENDING,
        triggered_by=None,  # System-triggered
        trigger_type=Run.TriggerType.SCHEDULED,
        code_snapshot=script.code,
    )

    # Update schedule tracking
    schedule.last_scheduled_run = timezone.now()
    schedule.save(update_fields=["last_scheduled_run"])

    # Queue for execution
    queue_script_run(run)

    # Update next_run cache
    schedule.next_run = ScheduleService._calculate_next_run(schedule)
    schedule.save(update_fields=["next_run"])

    logger.info(f"Scheduled run {run.id} created for script {script.name}")

    return {
        "success": True,
        "run_id": str(run.id),
        "script_id": script_id,
    }


def execute_package_operation(operation_id: str) -> dict:
    """
    Execute a package operation (install/uninstall) asynchronously.
    Called by django-q2 workers.

    Args:
        operation_id: The UUID of the PackageOperation record (as string)

    Returns:
        dict: Operation result summary
    """
    from core.services.environment_service import EnvironmentService

    try:
        operation = PackageOperation.objects.select_related("environment").get(
            id=UUID(operation_id)
        )
    except PackageOperation.DoesNotExist:
        logger.error(f"PackageOperation {operation_id} not found")
        return {"success": False, "error": "Operation not found"}
    except ValueError as e:
        logger.error(f"Invalid operation_id format: {operation_id} - {e}")
        return {"success": False, "error": f"Invalid UUID format: {e}"}

    # Update to running
    operation.status = PackageOperation.Status.RUNNING
    operation.started_at = timezone.now()
    operation.save(update_fields=["status", "started_at"])

    environment = operation.environment

    try:
        if operation.operation == PackageOperation.Operation.INSTALL:
            success, stdout, stderr = EnvironmentService.install_package(
                environment, operation.package_spec
            )
        elif operation.operation == PackageOperation.Operation.UNINSTALL:
            success, stdout, stderr = EnvironmentService.uninstall_package(
                environment, operation.package_spec
            )
        elif operation.operation == PackageOperation.Operation.BULK_INSTALL:
            success, stdout, stderr = EnvironmentService.install_requirements(
                environment, operation.package_spec
            )
        else:
            success, stdout, stderr = False, "", "Unknown operation type"

        operation.output = stdout
        operation.error = stderr
        operation.status = (
            PackageOperation.Status.SUCCESS
            if success
            else PackageOperation.Status.FAILED
        )

        # Update environment requirements cache on success
        if success:
            environment.requirements = EnvironmentService.pip_freeze(environment)
            environment.save(update_fields=["requirements", "updated_at"])

    except Exception as e:
        operation.status = PackageOperation.Status.FAILED
        operation.error = str(e)
        logger.exception(f"Package operation {operation_id} failed with exception")

    operation.completed_at = timezone.now()
    operation.save()

    logger.info(
        f"Package operation {operation_id} completed with status {operation.status}"
    )

    return {
        "success": operation.status == PackageOperation.Status.SUCCESS,
        "operation_id": operation_id,
        "status": operation.status,
    }


def cleanup_old_runs_task() -> dict:
    """
    Clean up old runs based on retention settings.
    Called by django-q2 scheduler or manually.

    Returns:
        dict: Cleanup result summary with count of deleted runs
    """
    from core.services.retention_service import RetentionService

    try:
        deleted_count = RetentionService.cleanup_all_runs()

        # Update last_cleanup_at timestamp
        settings = GlobalSettings.get_settings()
        settings.last_cleanup_at = timezone.now()
        settings.save(update_fields=["last_cleanup_at"])

        logger.info(f"Cleanup task completed: {deleted_count} runs deleted")
        return {"success": True, "deleted_count": deleted_count}
    except Exception as e:
        logger.exception("Cleanup task failed")
        return {"success": False, "error": str(e)}


def worker_heartbeat_task() -> dict:
    """
    Heartbeat task that runs periodically to indicate workers are alive.
    Updates the worker_heartbeat_at timestamp in GlobalSettings.

    This task is scheduled to run every minute when qcluster starts.
    It allows the system to detect when workers have stopped.

    Returns:
        dict: Result with timestamp of heartbeat
    """
    try:
        settings = GlobalSettings.get_settings()
        settings.worker_heartbeat_at = timezone.now()
        settings.save(update_fields=["worker_heartbeat_at"])

        logger.debug("Worker heartbeat updated")
        return {"success": True, "heartbeat_at": str(settings.worker_heartbeat_at)}
    except Exception as e:
        logger.exception("Worker heartbeat failed")
        return {"success": False, "error": str(e)}


def scheduled_backup_task() -> dict:
    """
    Execute a scheduled backup to S3.
    Called by django-q2 scheduler.

    Creates a backup, uploads to S3, and applies retention policy.

    Returns:
        dict: Backup result summary
    """
    from core.services.backup_service import BackupService
    from core.services.s3_service import S3Service
    from core.services.backup_schedule_service import BackupScheduleService

    settings = GlobalSettings.get_settings()

    # Validate S3 is configured
    if not settings.s3_enabled:
        logger.warning("Scheduled backup skipped - S3 not enabled")
        return {"success": False, "error": "S3 not enabled"}

    if not S3Service.is_configured():
        logger.warning("Scheduled backup skipped - S3 not configured")
        return {"success": False, "error": "S3 not configured"}

    try:
        # Create backup
        backup_data = BackupService.create_backup(
            include_runs=settings.s3_backup_include_runs,
            max_runs=settings.s3_backup_max_runs,
            include_package_operations=False,
            include_datastores=settings.s3_backup_include_datastores,
            created_by_user=None,  # System backup
        )

        # Serialize to gzip
        file_bytes, _ = BackupService.serialize_backup(
            backup_data,
            format=BackupService.FORMAT_GZIP,
        )

        # Generate S3 key
        key = S3Service.generate_backup_key()

        # Upload to S3
        result = S3Service.upload_file(file_bytes, key)

        if result["success"]:
            # Update tracking
            settings.s3_backup_last_run_at = timezone.now()
            settings.s3_backup_last_status = "success"
            settings.s3_backup_last_error = ""
            settings.s3_backup_last_size = result["size"]
            settings.save(update_fields=[
                "s3_backup_last_run_at",
                "s3_backup_last_status",
                "s3_backup_last_error",
                "s3_backup_last_size",
            ])

            # Apply retention cleanup
            deleted = BackupScheduleService.apply_retention()

            logger.info(f"Scheduled backup completed: {key} ({result['size']} bytes)")

            return {
                "success": True,
                "key": key,
                "size": result["size"],
                "retention_deleted": deleted,
            }
        else:
            # Update tracking with error
            settings.s3_backup_last_run_at = timezone.now()
            settings.s3_backup_last_status = "failed"
            settings.s3_backup_last_error = result.get("error", "Unknown error")
            settings.save(update_fields=[
                "s3_backup_last_run_at",
                "s3_backup_last_status",
                "s3_backup_last_error",
            ])

            logger.error(f"Scheduled backup failed: {result.get('error')}")
            return {"success": False, "error": result.get("error")}

    except Exception as e:
        # Update tracking with error
        settings.s3_backup_last_run_at = timezone.now()
        settings.s3_backup_last_status = "failed"
        settings.s3_backup_last_error = str(e)
        settings.save(update_fields=[
            "s3_backup_last_run_at",
            "s3_backup_last_status",
            "s3_backup_last_error",
        ])

        logger.exception("Scheduled backup task failed")
        return {"success": False, "error": str(e)}