from celery import shared_task
from backend.core.models import Organization, Document, Risk, Control, DocumentRiskControl
from backend.core.utils import get_top_risk, get_controls_for_risk
from django.shortcuts import get_object_or_404
from django.template import Template, Context
import yaml
from backend.core.models import DocumentTemplate
from .utils import send_payment_email

@shared_task
def create_document_for_organization(confirmation_email):
    organization = get_object_or_404(Organization, email=confirmation_email)

    top_risk_ids = get_top_risk(organization)
    top_risks = Risk.objects.filter(risk_id__in=top_risk_ids)

    template = DocumentTemplate.objects.first()
    template_content = template.content  

    risks_data = []
    for risk in top_risks:
        controls_data = []
        selected_controls = get_controls_for_risk(risk, organization)

        for control_id, weight in selected_controls:
            control = Control.objects.filter(id=control_id).first()
            if control:
                document = Document.objects.create(organization=organization)  
                DocumentRiskControl.objects.create(
                    document=document,
                    risk=risk,
                    control=control,
                    weight=weight,
                )

                controls_data.append({
                    'name': control.name,
                    'weight': weight,
                })

        risks_data.append({
            'name': risk.risk_name,
            'description': risk.description,  
            'controls': controls_data,
        })

    context = {
        'organization': organization,
        'risks': risks_data,
    }

    django_template = Template(template_content)
    rendered_content = django_template.render(Context(context))

    document = Document.objects.create(organization=organization, content=rendered_content)

    document.add_segment('title', f"{organization.name} Risk Assessment Report", 1)
    document.add_segment('introduction', f"This document outlines the risks and controls for {organization.name}.", 2)
    document.add_segment('body', rendered_content, 3)

    send_payment_email(confirmation_email)

    return document