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


class Command(BaseCommand):
    help = "Export risks to a CSV file"

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

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

        with open(csv_file_path, mode="w", newline="") as csv_file:
            fieldnames = [
                'Risk ID','Category','Risk Name', 'Primary Impact', 'Secondary Impact',
                'Tertiary Impact', 'Detection Difficulty', 'Recovery Complexity',
                'Business Impact Severity'

                ]
            writer = csv.DictWriter(csv_file, fieldnames=fieldnames)

            writer.writeheader()
            for risk in Risk.objects.all():
                writer.writerow({
                    'Risk ID': risk.risk_id,
                    'Category': risk.category,
                    'Risk Name': risk.risk_name,
                    'Primary Impact': risk.primary_impact,
                    'Secondary Impact': risk.secondary_impact,
                    'Tertiary Impact': risk.tretiary_impact,
                    'Detection Difficulty': risk.detection_difficulty,
                    'Recovery Complexity': risk.recovery_complexity,
                    'Business Impact Severity': risk.businnes_impact_severity,

                })
        self.stdout.write(self.style.SUCCESS('Risks exported successfully'))