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

class Command(BaseCommand):
    help = "Export controls 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 = [
                "Risk #",
                "Risk Description",
                "CIS v8.1 Safeguards (Sub-Controls)",
                "Weight (0-10)",
            ]
            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)

            writer.writeheader()

            for control in Control.objects.select_related("risk").all():
                writer.writerow({
                    "Risk #": control.risk.risk_id,
                    "Risk Description": control.risk.risk_name,
                    "CIS v8.1 Safeguards (Sub-Controls)": control.safeguard,
                    "Weight (0-10)": control.weight,
                })

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