import csv
from django.core.management.base import BaseCommand
from backend.core.models import Control

class Command(BaseCommand):
    help = "Export controls (only safeguard) to CSV file"

    def add_arguments(self, parser):
        parser.add_argument("csv_file", type=str, help="CSV file to export controls to")

    def handle(self, *args, **options):
        csv_file_path = options["csv_file"]

        with open(csv_file_path, mode="w", newline="", encoding="utf-8") as csv_file:
            fieldnames = ["CIS v8.1 Safeguards (Sub-Controls)"] 
            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)

            writer.writeheader()

            for control in Control.objects.all():
                writer.writerow({
                    "CIS v8.1 Safeguards (Sub-Controls)": control.name,
                })

        self.stdout.write(self.style.SUCCESS(f"Controls exported successfully to {csv_file_path}"))
