"""Batch enrichment used by document read serializers."""

from django.contrib.auth import get_user_model
from django.contrib.contenttypes.models import ContentType
from django.db.models import Count

from communications.models import Message
from documents.models import Document, DocumentAcknowledgement
from employees.models import EmployeeDepartment, RoleAssignment

User = get_user_model()


def attach_linked_task_payloads(documents: list[Document], *, user) -> None:
    """Attach safe task badges to a materialized document collection."""
    if not documents:
        return

    from tasks.models import TaskLinkedObjectKind
    from tasks.presentation import linked_task_badges_by_object

    mapping = linked_task_badges_by_object(
        model=Document,
        kind=TaskLinkedObjectKind.DOCUMENT,
        object_ids=[document.id for document in documents],
        user=user,
    )

    for document in documents:
        document._linked_task_payloads = mapping.get(document.id, [])


def _department_ids(documents: list[Document]) -> set[int]:
    return {
        department.id
        for document in documents
        for department in (
            list(document.departments.all())
            + list(document.acknowledgement_departments.all())
        )
    }


def _employees_by_department(department_ids: set[int]) -> dict[int, set[int]]:
    mapping = {department_id: set() for department_id in department_ids}
    if not department_ids:
        return mapping

    membership_rows = EmployeeDepartment.objects.filter(
        department_id__in=department_ids,
        is_active=True,
        employee__is_active=True,
    ).values_list("department_id", "employee_id")
    role_rows = RoleAssignment.objects.filter(
        role__department_id__in=department_ids,
        is_active=True,
        employee__is_active=True,
    ).values_list("role__department_id", "employee_id")
    for department_id, employee_id in (*membership_rows, *role_rows):
        mapping.setdefault(department_id, set()).add(employee_id)
    return mapping


def _acknowledged_users(document_ids: list[int]) -> dict[int, set[int]]:
    mapping = {document_id: set() for document_id in document_ids}
    rows = DocumentAcknowledgement.objects.filter(
        document_id__in=document_ids,
        user__is_active=True,
    ).values_list("document_id", "user_id")
    for document_id, user_id in rows:
        mapping.setdefault(document_id, set()).add(user_id)
    return mapping


def _comment_counts(document_ids: list[int]) -> dict[int, int]:
    document_content_type = ContentType.objects.get_for_model(Document)
    return {
        row["chat__context_object_id"]: row["total"]
        for row in Message.objects.filter(
            chat__type="comments",
            chat__context_content_type=document_content_type,
            chat__context_object_id__in=document_ids,
            is_deleted=False,
        )
        .values("chat__context_object_id")
        .annotate(total=Count("id"))
    }


def _document_access_ids(
    document: Document,
    active_user_ids: set[int],
    employees_by_department: dict[int, set[int]],
) -> set[int]:
    if document.sent_to_all:
        return set(active_user_ids)
    access_ids = {
        recipient.id
        for recipient in document.recipients.all()
        if recipient.id in active_user_ids
    }
    for department in document.departments.all():
        access_ids.update(employees_by_department.get(department.id, set()))
    return access_ids


def _acknowledgement_ids(
    document: Document,
    access_ids: set[int],
    active_user_ids: set[int],
    employees_by_department: dict[int, set[int]],
) -> set[int]:
    if not document.acknowledgement_required:
        return set()
    if document.acknowledgement_for_all:
        return set(access_ids)

    required_ids = {
        recipient.id
        for recipient in document.acknowledgement_recipients.all()
        if recipient.id in active_user_ids
    }
    for department in document.acknowledgement_departments.all():
        required_ids.update(employees_by_department.get(department.id, set()))
    return required_ids & access_ids


def attach_document_metrics(documents: list[Document]) -> None:
    """Attach comment and acknowledgement aggregates without N+1 queries."""
    if not documents:
        return

    document_ids = [document.id for document in documents]
    active_user_ids = set(
        User.objects.filter(is_active=True).values_list("id", flat=True)
    )
    employees_by_department = _employees_by_department(_department_ids(documents))
    acknowledged_users = _acknowledged_users(document_ids)
    comment_counts = _comment_counts(document_ids)

    for document in documents:
        access_ids = _document_access_ids(
            document,
            active_user_ids,
            employees_by_department,
        )
        required_ids = _acknowledgement_ids(
            document,
            access_ids,
            active_user_ids,
            employees_by_department,
        )
        document._acknowledgement_user_ids = required_ids
        document._acknowledgement_total = len(required_ids)
        document._acknowledged_count = len(
            required_ids & acknowledged_users.get(document.id, set())
        )
        document._comments_count = comment_counts.get(document.id, 0)