import os
import csv
import tempfile
from django.core.management import call_command
from django.test import TestCase
from backend.core.models import Risk, Control

class ImportControlsCommandTest(TestCase):
    
    def setUp(self):
        """Create a temporary CSV file with test control data"""
        self.temp_csv = tempfile.NamedTemporaryFile(mode="w", delete=False, newline='', encoding='utf-8')
        fieldnames = [
            "Risk #", "Risk Description", "CIS v8.1 Safeguards (Sub-Controls)", "Weight (0-10)"
        ]
        
        writer = csv.DictWriter(self.temp_csv, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerow({
            "Risk #": "1",
            "Risk Description": "Test Risk 1",
            "CIS v8.1 Safeguards (Sub-Controls)": "Test Safeguard 1",
            "Weight (0-10)": "5"
        })
        writer.writerow({
            "Risk #": "2",
            "Risk Description": "Test Risk 2",
            "CIS v8.1 Safeguards (Sub-Controls)": "Test Safeguard 2",
            "Weight (0-10)": "10"
        })
        self.temp_csv.close()

    def test_import_controls_command(self):
        """Test importing controls from a CSV file"""
        self.assertEqual(Control.objects.count(), 0)
        
        self.assertEqual(Risk.objects.count(), 0)

        Risk.objects.create(risk_id=1, risk_name="Test Risk 1")
        Risk.objects.create(risk_id=2, risk_name="Test Risk 2")

        call_command('import_controls', self.temp_csv.name)
        
        self.assertEqual(Control.objects.count(), 2)

        control1 = Control.objects.get(risk__risk_id=1)
        self.assertEqual(control1.safeguard, "Test Safeguard 1")
        self.assertEqual(control1.weight, 5)

        control2 = Control.objects.get(risk__risk_id=2)
        self.assertEqual(control2.safeguard, "Test Safeguard 2")
        self.assertEqual(control2.weight, 10)

    def tearDown(self):
        """Remove temporary CSV file after test"""
        os.remove(self.temp_csv.name)