from django.core.management.base import BaseCommand
from tqdm import tqdm

from users.models import Task


class Command(BaseCommand):
    help = "Manage tasks"

    def add_arguments(self, parser):
        parser.add_argument("--id", action="append")
        parser.add_argument("--user", action="append")
        parser.add_argument("--type", action="append")
        parser.add_argument("--pending", action="store_true")
        parser.add_argument("--failed", action="store_true")
        parser.add_argument("--complete", action="store_true")
        parser.add_argument("--list", action="store_true")
        parser.add_argument("--prune", action="store_true")
        parser.add_argument("--rerun", action="store_true")
        parser.add_argument("--requeue", action="store_true")
        # parser.add_argument("--set-fail", action="store_true")
        parser.add_argument("--delete", action="store_true")

    def handle(self, *args, **options):
        tasks = Task.objects.all().order_by("id")
        states = []
        if options["pending"]:
            states += [Task.States.pending]
        if options["failed"]:
            states += [Task.States.failed]
        if options["complete"]:
            states += [Task.States.complete]
        if states:
            tasks = tasks.filter(state__in=states)
        if options["user"]:
            tasks = tasks.filter(user__username__in=options["user"])
        if options["id"]:
            tasks = tasks.filter(id__in=options["id"])
        if options["type"]:
            tasks = tasks.filter(type__in=options["type"])
        if options["list"]:
            for task in tasks.order_by("-created_time"):
                self.stdout.write(
                    str(task.pk).ljust(10)
                    + str(task.type).ljust(30)
                    + str(Task.States(task.state).label).ljust(10)
                    + task.created_time.strftime("%Y-%m-%d %H:%M  ")
                    + task.edited_time.strftime("%Y-%m-%d %H:%M  ")
                    + str(task.user)
                )
        if options["rerun"]:
            for task in tqdm(tasks):
                task.state = Task.States.pending
                task.save(update_fields=["state"])
                Task._execute(task.pk)
        if options["requeue"]:
            for task in tqdm(tasks):
                task.state = Task.States.pending
                task.save(update_fields=["state"])
                task.enqueue()
        if options["delete"]:
            for task in tqdm(tasks):
                task.delete()