from math import ceil
from django.core.management.base import BaseCommand

from kitsune.search.v2.es7_utils import get_doc_types, index_objects_bulk


class Command(BaseCommand):
    help = "Reindex ES7 documents"

    def add_arguments(self, parser):
        parser.add_argument(
            "--limit",
            type=str,
            dest="limit",
            nargs="*",
            default="",
            help="Limit to specific doc types",
        )
        parser.add_argument(
            "--percentage", type=float, default=100,
        )
        parser.add_argument(
            "--bulk-count", type=int, default=100,
        )

    def handle(self, *args, **kwargs):
        doc_types = get_doc_types()

        limit = kwargs["limit"]
        if limit:
            doc_types = [dt for dt in doc_types if dt.__name__ in limit]

        progress_msg = "Indexed {progress} out of {count}"

        for dt in doc_types:
            self.stdout.write("Reindexing: {}".format(dt.__name__))

            model = dt.get_model()
            qs = model.objects.all()
            count = qs.count()

            percentage = kwargs["percentage"]
            total = count
            if percentage < 100:
                count = int(count * percentage / 100)
                qs = qs[:count]
            print("Indexing {}%, so {} documents out of {}".format(percentage, count, total))

            id_list = list(qs.values_list("pk", flat=True))
            bulk_count = kwargs["bulk_count"]

            for x in range(ceil(count / bulk_count)):
                start = x * bulk_count
                end = start + bulk_count
                index_objects_bulk.delay(dt.__name__, id_list[start:end])
                print(progress_msg.format(progress=min(end, count), count=count))

@@ -535,97 +535,555 @@
 100,
-%0A        )%0A        parser.add_argument(%0A            %22--bulk-count%22, type=int, default=100
+ help=%22Index a percentage of total documents%22,%0A        )%0A        parser.add_argument(%0A            %22--count%22,%0A            type=int,%0A            default=None,%0A            help=%22Index a set number of documents per type (overrides --percentage)%22,%0A        )%0A        parser.add_argument(%0A            %22--bulk-count%22, type=int, default=100, help=%22Index this number of documents at once%22,%0A        )%0A        parser.add_argument(%0A            %22--print-sql-count%22,%0A            action=%22store_true%22,%0A            help=%22Print the number of SQL statements executed%22
 ,%0A  
@@ -1351,16 +1351,115 @@
 ount%7D%22%0A%0A
+        if kwargs%5B%22print_sql_count%22%5D:%0A            from django.db import connection, reset_queries%0A%0A
         
@@ -1720,24 +1720,165 @@
 tal = count%0A
+            if count := kwargs%5B%22count%22%5D:%0A                print(%22Indexing %7B%7D documents out of %7B%7D%22.format(count, total))%0A            else:%0A    
             
@@ -1898,16 +1898,20 @@
  %3C 100:%0A
+    
         
@@ -1968,24 +1968,28 @@
             
+    
 qs = qs%5B:cou
@@ -1992,16 +1992,20 @@
 :count%5D%0A
+    
         
@@ -2406,16 +2406,183 @@
 t:end%5D)%0A
+                if kwargs%5B%22print_sql_count%22%5D:%0A                    print(%22%7B%7D SQL queries executed%22.format(len(connection.queries)))%0A                    reset_queries()%0A
         

add count and print-sql-count args to es7_reindex command