import os
import csv
from django.core.management import call_command
from django.test import TestCase
from django.conf import settings
from backend.core.models import Risk

class ExportRisksCommandTest(TestCase):

    def setUp(self):
        self.risk = Risk.objects.create(
            risk_id="1",
            category="Security",
            risk_name="Data Breach",
            primary_impact="High",
            secondary_impact="Medium",
            tretiary_impact="Low",
            detection_difficulty="High",
            recovery_complexity="Medium",
            businnes_impact_severity="High"
        )
        self.csv_file_path = os.path.join(settings.BASE_DIR, "test_risks.csv")

    def tearDown(self):
        if os.path.exists(self.csv_file_path):
            os.remove(self.csv_file_path)

    def test_export_risks_to_csv(self):
        call_command('export_risks', self.csv_file_path)

        self.assertTrue(os.path.exists(self.csv_file_path))

        with open(self.csv_file_path, mode='r') as csv_file:
            reader = csv.DictReader(csv_file)
            rows = list(reader)

            self.assertEqual(len(rows), 1)

            row = rows[0]
            self.assertEqual(row['Risk ID'], self.risk.risk_id)
            self.assertEqual(row['Category'], self.risk.category)
            self.assertEqual(row['Risk Name'], self.risk.risk_name)
            self.assertEqual(row['Primary Impact'], self.risk.primary_impact)
            self.assertEqual(row['Secondary Impact'], self.risk.secondary_impact)
            self.assertEqual(row['Tertiary Impact'], self.risk.tretiary_impact)
            self.assertEqual(row['Detection Difficulty'], self.risk.detection_difficulty)
            self.assertEqual(row['Recovery Complexity'], self.risk.recovery_complexity)
            self.assertEqual(row['Business Impact Severity'], self.risk.businnes_impact_severity)

    def test_no_risks_in_db(self):
        Risk.objects.all().delete()

        call_command('export_risks', self.csv_file_path)

        self.assertTrue(os.path.exists(self.csv_file_path))

        with open(self.csv_file_path, mode='r') as csv_file:
            reader = csv.DictReader(csv_file)
            rows = list(reader)

            self.assertEqual(len(rows), 0)
