from django.test import TestCase
from unittest.mock import patch, MagicMock
from django.contrib.staticfiles.testing import StaticLiveServerTestCase
from io import BytesIO
import base64
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt

from backend.core.models import Organization, Risk, Control, Document
from backend.core.utils import *

class UtilsTests(TestCase):
    def setUp(self):
        self.organization = Organization.objects.create(
            id=1,
            name="Test Organization",
            email="test@example.com",
            employee_headcount="100-500",
            annual_revenue="$1M-$10M",
            critical_applications="5-10",
            compliance_frameworks=["Ab", "Ba"],
            industry_sector="Technology",
            it_dependency=8,
            data_sensitivity="High",
            network_infrastructure="Cloud-based",
            remote_workforce_percentage="50%",
            third_party_vendor_access="10-20",
            internal_software_development="Moderate",
            geographic_scope="Global",
            customer_base="Enterprise",
            customer_type="B2B",
            product_portfolio="Diverse",
            supplier_base="International",
            it_infrastructure=["Cloud", "On-Premise"],
            intellectual_property=["Patents", "Trademarks"],
            sensitive_data=["PII", "Financial Data"],
            integration_level="Highly Integrated"
        )

        self.risk = Risk.objects.create(
            risk_id=1,
            risk_name="Test Risk",
            category="Security",
            primary_impact="Financial"
        )

        self.controls = [Control.objects.create(id=i, name=f"Control {i}") for i in range(1, 11)]

    def test_extract_organization_details(self):
        details = extract_organization_details(self.organization)
        self.assertNotIn('name', details)
        self.assertNotIn('email', details)
        self.assertIn("What is your organization's current employee headcount?", details)
        self.assertEqual(details["What is your organization's current employee headcount?"], "100-500")

    @patch('backend.core.utils.OpenAI')
    def test_get_top_risk(self, mock_openai):
        mock_client = MagicMock()
        mock_openai.return_value = mock_client
        mock_response = MagicMock()
        mock_response.choices[0].message.content = "1,2,3"
        mock_client.chat.completions.create.return_value = mock_response

        risks = get_top_risk(self.organization)
        self.assertEqual(risks, [1,2,3])

    @patch('backend.core.utils.OpenAI')
    def test_get_controls_for_risk(self, mock_openai):
        mock_client = MagicMock()
        mock_openai.return_value = mock_client
        mock_response = MagicMock()
        
        control_lines = [f"{i} : 4 : 3" for i in range(1, 11)]
        mock_response.choices[0].message.content = "\n".join(control_lines)
        mock_client.chat.completions.create.return_value = mock_response

        controls = get_controls_for_risk(self.risk, self.organization)
        self.assertEqual(len(controls), 10)
        self.assertEqual(controls[0][0], 1)

    @patch('backend.core.utils.HTML')
    def test_generate_pdf(self, mock_html):
        mock_instance = MagicMock()
        mock_instance.write_pdf.return_value = b'PDF_CONTENT'
        mock_html.return_value = mock_instance

        doc = Document.objects.create(organization=self.organization)
        response = generate_pdf(doc)
        
        self.assertEqual(response.status_code, 200)
        self.assertEqual(response['Content-Type'], 'application/pdf')

    def test_calculate_aggregate_weight(self):
        controls = [{'weight': 5}, {'weight': 3}]
        self.assertEqual(calculate_aggregate_weight(controls), 8)

    def test_calculate_aggregate_likelihood(self):
        controls = [{'likelihood': 2}, {'likelihood': 4}]
        self.assertEqual(calculate_aggregate_likelihood(controls), 6)

    def test_map_weight_to_impact_likelihood(self):
        impact, likelihood = map_weight_to_impact_likelihood(50, 30, 50)
        self.assertAlmostEqual(impact, 5.0)
        self.assertAlmostEqual(likelihood, 3.0)

    @patch('pdf2image.convert_from_bytes')
    @patch('backend.core.utils.HTML')
    def test_generate_first_page_image(self, mock_html, mock_convert):
        mock_pdf_instance = MagicMock()
        mock_pdf_instance.write_pdf.return_value = b'PDF_CONTENT'
        mock_html.return_value = mock_pdf_instance
        
        mock_image = MagicMock()
        mock_convert.return_value = [mock_image]

        doc = Document.objects.create(organization=self.organization)
        img_io = generate_first_page_image(doc)
        
        self.assertIsInstance(img_io, BytesIO)
        mock_convert.assert_called_once_with(b'PDF_CONTENT', first_page=1, last_page=1)

    def test_generate_risk_graph(self):
        risks_with_controls = [
            {'risk': {'id': 1}, 'impact': 5.0, 'likelihood': 3.0},
            {'risk': {'id': 2}, 'impact': 4.0, 'likelihood': 4.0}
        ]
        
        graph_data = generate_risk_graph(risks_with_controls)
        self.assertIsInstance(graph_data, str)
        self.assertTrue(len(graph_data) > 1000)