from backend.core.models import DocumentRiskControl
from backend.core.utils import calculate_aggregate_likelihood, calculate_aggregate_weight, map_weight_to_impact_likelihood


def risk_matrix_table():
    likelihood_labels = [
        "Almost Certain (90-100%) (5)",
        "Probable (51-89%) (4)",
        "Possible (25-50%) (3)",
        "Unlikely (11-24%) (2)",
        "Rare (0-10%) (1)"
    ]

    impact_labels = [
        "Insignificant (1)",
        "Significant (2)",
        "Severe (3)",
        "Material (4)",
        "Major (5)"
    ]

    color_mapping = {
        "Very Low": "lightgreen",
        "Low": "green",
        "Medium": "yellow",
        "High": "orange",
        "Critical": "red"
    }

    def get_label(score):
        if score <= 2:
            return "Very Low"
        elif score <= 4:
            return "Low"
        elif score <= 10:
            return "Medium"
        elif score <= 16:
            return "High"
        else:
            return "Critical"

    table_matrix_risk = [["Likelihood ↓ / Impact →"] + impact_labels]

    for likelihood in range(5, 0, -1):
        row = [likelihood_labels[5 - likelihood]]
        for impact in range(1, 6):
            score = likelihood * impact
            label = get_label(score)
            color_class = color_mapping[label]
            row.append((score, label, color_class))
        table_matrix_risk.append(row)

    return table_matrix_risk

def get_risk_table(document):
    risks = (
        DocumentRiskControl.objects
        .filter(document=document)
        .values('risk', 'risk__risk_name')
        .distinct()
    )

    risks_with_controls = []

    for risk_entry in risks:
        risk = {
            'id': risk_entry['risk'],
            'name': risk_entry['risk__risk_name']
        }

        controls = (
            DocumentRiskControl.objects
            .filter(document=document, risk_id=risk['id'])
            .values('control', 'control__name', 'weight', 'likelihood')
            .distinct()
        )
        max_weight = 10*5
        total_weight = calculate_aggregate_weight(controls)
        total_likelihood = calculate_aggregate_likelihood(controls)
        impact, likelihood = map_weight_to_impact_likelihood(total_weight, total_likelihood, max_weight)
        r_impact = round(impact)
        r_likelihood = round(likelihood)
        residua_impact = r_impact - 1 if r_impact > 2 else r_impact
        residual_likelihood = r_likelihood - 1 if r_likelihood > 2 else r_likelihood
        risks_with_controls.append({
            'risk': risk,
            'controls': list(controls),
            'total_weight': total_weight,
            'impact': impact,
            'likelihood': likelihood,
            'r_impact': r_impact,
            'r_likelihood': r_likelihood,
            'risk_score': r_impact * r_likelihood,
            'residual_impact': residua_impact,
            'residual_likelihood': residual_likelihood,
            'residual_risk_score': residua_impact * residual_likelihood,
        })
    risks_with_controls.sort(key=lambda x: x['risk_score'], reverse=True)

    return risks_with_controls