from django.core.management.base import BaseCommand
from django.utils import timezone
from tasks.models import Task

class Command(BaseCommand):
    help = 'Check and update tasks that have timed out'

    def handle(self, *args, **kwargs):
        now = timezone.now()
        
        # Get all tasks that are in progress (assigned, running, retrying)
        in_progress_tasks = Task.objects.filter(
            status__in=['assigned', 'running', 'retrying']
        )
        
        timeout_count = 0
        
        for task in in_progress_tasks:
            # Calculate when the task should timeout
            if task.started_at:
                # If task has started, calculate from started_at
                timeout_time = task.started_at + timezone.timedelta(seconds=task.timeout_seconds)
            else:
                # If task hasn't started, calculate from created_at
                timeout_time = task.created_at + timezone.timedelta(seconds=task.timeout_seconds)
            
            # Check if task has timed out
            if now > timeout_time:
                task.status = 'timeout'
                task.save()
                timeout_count += 1
                self.stdout.write(self.style.WARNING(f'Task {task.id} "{task.name}" has timed out'))
        
        if timeout_count > 0:
            self.stdout.write(self.style.SUCCESS(f'Successfully updated {timeout_count} timed out tasks'))
        else:
            self.stdout.write(self.style.INFO('No tasks have timed out'))