from django.core.management.base import BaseCommand

from candidates.models import Ballot
from utils.dict_io import BufferDictWriter


class Command(BaseCommand):

    help = """
    Create a CSV file that can be populated and imported by
    candidates_import_statements_of_persons_nominated
    """

    fieldnames = [
        "ballot_paper_id",
        "Election name",
        "Area name",
        "Council webpage likely to link to SOPN",
        "Link to PDF",
        "Notes",
    ]

    def add_arguments(self, parser):
        parser.add_argument("--election-date", action="store", required=True)

        parser.add_argument(
            "--non-current",
            action="store_true",
            help="Also include elections marked as not current",
        )

    def handle(self, *args, **options):

        out_csv = BufferDictWriter(self.fieldnames)
        out_csv.writeheader()

        qs = (
            Ballot.objects.filter(
                election__election_date=options["election_date"]
            )
            .select_related("election", "post")
            .order_by("election__slug")
        )

        if not options["non_current"]:
            qs = qs.filter(election__current=True)

        for ballot in qs:
            row = {
                "ballot_paper_id": ballot.ballot_paper_id,
                "Election name": ballot.election.name,
                "Area name": ballot.post.label,
                "Council webpage likely to link to SOPN": "",
                "Link to PDF": "",
                "Notes": "",
            }
            out_csv.writerow(row)

        self.stdout.write(out_csv.output)

@@ -46,16 +46,49 @@
 ommand%0A%0A
+from django.conf import settings%0A
 from can
@@ -116,16 +116,16 @@
  Ballot%0A
-
 from uti
@@ -796,24 +796,187 @@
 %0A        )%0A%0A
+        parser.add_argument(%0A            %22--nations%22,%0A            action=%22store%22,%0A            help=%22Limit ballots to given nations. List of %5BE,N,S,W%5D%22,%0A        )%0A%0A
     def hand
@@ -1354,16 +1354,16 @@
 rent%22%5D:%0A
-
         
@@ -1398,32 +1398,447 @@
 _current=True)%0A%0A
+        if options%5B%22nations%22%5D:%0A            nations = set(%0A                %5Bn.strip().upper() for n in options%5B%22nations%22%5D.split(%22,%22)%5D%0A            )%0A            if not nations.issubset(set(settings.NUTS_TO_NATION.keys())):%0A                raise ValueError(%0A                    f%22Nations must be in %7Bset(settings.NUTS_TO_NATION.keys())%7D (%7Bnations%7D provided)%22%0A                )%0A            qs = qs.by_nation(*nations)%0A%0A
         for ball

Allow filtering SOPN sheet by nations