"""Tests for risk aggregation utilities."""

import pytest

from sentinel.models import RiskScore
from sentinel.risk_aggregation import (
    format_scores_for_llm,
    format_scores_for_pdf,
    group_scores_by_cancer_type,
)


def test_group_scores_by_cancer_type():
    """Test grouping risk scores by cancer type."""
    scores = [
        RiskScore(
            name="Gail Model",
            score="5%",
            cancer_type="Breast Cancer",
            description="5-year risk",
        ),
        RiskScore(
            name="Claus Model",
            score="3%",
            cancer_type="Breast Cancer",
            description="Lifetime risk",
        ),
        RiskScore(
            name="PLCOm2012",
            score="2%",
            cancer_type="Lung Cancer",
            description="6-year risk",
        ),
    ]

    grouped = group_scores_by_cancer_type(scores)

    assert len(grouped) == 2
    assert "Breast Cancer" in grouped
    assert "Lung Cancer" in grouped
    assert len(grouped["Breast Cancer"]) == 2
    assert len(grouped["Lung Cancer"]) == 1
    assert grouped["Breast Cancer"][0].name == "Gail Model"
    assert grouped["Breast Cancer"][1].name == "Claus Model"


def test_group_scores_empty():
    """Test grouping with empty list."""
    grouped = group_scores_by_cancer_type([])
    assert grouped == {}


def test_group_scores_no_cancer_type():
    """Test grouping with scores that have no cancer type."""

    scores = [
        RiskScore(name="Test Model", score="5%", cancer_type=None),
    ]

    with pytest.raises(ValueError, match=r"Test Model.*missing cancer_type"):
        group_scores_by_cancer_type(scores)


def test_format_scores_for_llm():
    """Test formatting scores for LLM context."""
    scores = [
        RiskScore(
            name="Gail Model",
            score="5%",
            cancer_type="Breast Cancer",
            description="5-year risk",
            interpretation="Low to moderate risk",
            references=["Gail et al., 1989"],
        ),
    ]

    grouped = group_scores_by_cancer_type(scores)
    formatted = format_scores_for_llm(grouped)

    assert "# Calculated Risk Scores (Ground Truth)" in formatted
    assert "Breast Cancer" in formatted
    assert "Gail Model" in formatted
    assert "5%" in formatted
    assert "5-year risk" in formatted
    assert "Low to moderate risk" in formatted
    assert "Gail et al., 1989" in formatted
    assert "DO NOT generate your own risk levels" in formatted


def test_format_scores_for_llm_empty():
    """Test formatting empty scores for LLM."""
    formatted = format_scores_for_llm({})
    assert formatted == "No risk scores calculated."


def test_format_scores_for_pdf():
    """Test formatting scores for PDF presentation."""
    scores = [
        RiskScore(
            name="Gail Model",
            score="5%",
            cancer_type="Breast Cancer",
        ),
        RiskScore(
            name="PLCOm2012",
            score="2%",
            cancer_type="Lung Cancer",
        ),
    ]

    grouped = group_scores_by_cancer_type(scores)
    formatted = format_scores_for_pdf(grouped)

    assert len(formatted) == 2
    assert formatted[0][0] == "Breast Cancer"
    assert formatted[1][0] == "Lung Cancer"
    assert len(formatted[0][1]) == 1
    assert len(formatted[1][1]) == 1