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

class Command(BaseCommand):
    help = "Import controls from CSV file"

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

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

        with open(csv_file_path, mode="r", encoding="utf-8") as csv_file:
            reader = csv.DictReader(csv_file)

            for row in reader:
                risk_id = int(row["Risk #"].strip())  
                risk_desc = row["Risk Description"].strip() 
                safeguard = row["CIS v8.1 Safeguards (Sub-Controls)"].strip()
                
                weight = int(row["Weight (0-10)"].strip()) if row["Weight (0-10)"].strip().isdigit() else 0 

                risk, created = Risk.objects.get_or_create(
                    risk_id=risk_id,
                    defaults={"risk_name": risk_desc},
                )

                Control.objects.update_or_create(
                    risk=risk,
                    safeguard=safeguard,
                    defaults={
                        "description": risk_desc,
                        "weight": weight,
                    },
                )

        self.stdout.write(self.style.SUCCESS("Controls imported successfully!"))